adolf/router.py

import asyncio
import re
import math
from typing import Optional
from openai import AsyncOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from fast_tools import FastToolRunner

# ── Regex pre-classifiers ─────────────────────────────────────────────────────

# Complex: keyword triggers that reliably signal deep multi-source research
_COMPLEX_PATTERNS = re.compile(
    r"(?:^|\s)("
    r"research|investigate|deep.dive|think carefully"
    r"|write a (?:detailed|comprehensive|full|thorough|complete)"
    r"|compare all|find and (?:compare|summarize|analyze)"
    r"|in[- ]depth analysis|comprehensive guide"
    r"|detailed (?:report|analysis|comparison|breakdown|overview)"
    r"|everything about|all (?:major|available|self-hosted|open.source)"
    r"|pros and cons|with (?:sources|citations|references)"
    # Russian complex research keywords (no trailing \b — stems like подробн match подробное/подробный)
    r"|исследуй|изучи все|сравни все|найди и сравни|найди и опиши"
    r"|напиши подробн|напиши детальн|напиши полн"
    r"|подробный отчет|детальн\w+ (?:анализ|сравнение|отчет)"
    r"|подробное (?:руководство|сравнение)|полное руководство"
    r"|все варианты|все способы|все доступные|все самохостируемые|все платформы"
    r"|лучшие практики|все инструменты|все решения|все протоколы"
    r"|найди детальн|найди и кратко опиши"
    r"|изучи свежие|изучи лучши|изучи все"
    r"|сравни все\b"
    r")",
    re.IGNORECASE,
)

# Light: trivial queries that need no tools or memory
_LIGHT_PATTERNS = re.compile(
    r"^("
    # Greetings / farewells
    r"hi|hello|hey|yo|sup|howdy|good morning|good evening|good night|good afternoon"
    r"|bye|goodbye|see you|cya|later|ttyl"
    # Acknowledgements / small talk
    r"|thanks?|thank you|thx|ty|ok|okay|k|cool|great|awesome|perfect|sounds good|got it|nice|sure"
    r"|how are you|how are you\?|how are you doing(\s+today)?[?!.]*"
    r"|what.?s up"
    # Calendar facts
    r"|what\s+day\s+(comes\s+after|follows|is\s+after)\s+\w+[?!.]*"
    r"|what\s+comes\s+after\s+\w+[?!.]*"
    # Acronym expansions
    r"|what\s+does\s+\w+\s+stand\s+for[?!.]*"
    # Russian greetings / farewells / acknowledgements
    r"|привет|пока|спасибо|здравствуй|здравствуйте|добрый день|добрый вечер|доброе утро"
    r"|окей|хорошо|отлично|понятно|ок|ладно|договорились|спс|благодарю"
    r"|пожалуйста|не за что|всё понятно|ясно"
    r"|как дела|как ты|как жизнь|всё хорошо|всё ок"
    # Russian tech definitions — static knowledge (no tools needed)
    r"|что\s+такое\s+\S+"
    r"|что\s+означает\s+\S+"
    r"|сколько\s+(?:бит|байт|байтов|мегабайт|мегабайтов|гигабайт|гигабайтов)(?:\s+\w+)*"
    # Compound Russian greetings
    r"|привет[,!]?\s+как\s+дела"
    r"|добрый\s+(?:день|вечер|утро)[,!]?\s+как\s+дела"
    r")[\s!.?]*$",
    re.IGNORECASE,
)

# ── Semantic router utterances ────────────────────────────────────────────────
# These are embedded at startup. New messages are classified by cosine
# similarity — whichever tier's centroid is closest wins.
_LIGHT_UTTERANCES = [
    # General facts (English)
    "what is 2+2",
    "what is the capital of France",
    "name the three primary colors",
    "tell me a short joke",
    "is the sky blue",
    "is water wet",
    "how many days in a week",
    "what is the speed of light",
    "what is the boiling point of water",
    "spell the word beautiful",
    "what color is the ocean",
    "how many inches in a foot",
    "who wrote hamlet",
    "what is pi",
    "what year did world war two end",
    "what is the largest planet",
    "how many continents are there",
    "what does DNA stand for",
    "what language do they speak in Brazil",
    "what is the square root of 144",
    # Tech definitions — static knowledge (English)
    "what is Docker",
    "what is a VPN",
    "what is SSH",
    "what is a reverse proxy",
    "what is an API",
    "what is a firewall",
    "what is a container",
    "what is DNS",
    "what is HTTPS",
    "what is a load balancer",
    "what is Kubernetes",
    "what is Git",
    "what is a network port",
    "what is an IP address",
    "what is a subnet mask",
    "what is the OSI model",
    "how many bits in a byte",
    "how many bytes in a gigabyte",
    "what is TCP",
    "what is a REST API",
    # Russian — static facts and definitions
    "что такое IP-адрес",
    "что такое VPN",
    "что такое Docker",
    "что такое DNS",
    "что такое SSH",
    "что означает API",
    "сколько байт в гигабайте",
    "сколько бит в байте",
    "что такое Zigbee",
    "что такое Z-Wave",
    "что такое брандмауэр",
    "что такое виртуальная машина",
    "что такое обратный прокси",
    "привет",
    "пока",
    "спасибо",
    "как дела",
    "что такое Matter протокол",
    "сколько планет в солнечной системе",
    "чему равно число Пи",
    # Russian — more static definitions
    "что такое TCP/IP",
    "что такое подсеть",
    "скорость света",
    "сколько дней в году",
    "что такое Kubernetes",
    "что такое Git",
    "что такое REST API",
    "что такое TCP",
    "что такое UDP",
    "что такое VLAN",
    "сколько мегабайт в гигабайте",
    "что такое процессор",
    "что такое оперативная память",
    "что такое виртуализация",
    "что такое Linux",
    "что такое умный дом",
    "что такое Home Assistant",
    "что такое Matter",
]

_MEDIUM_UTTERANCES = [
    # English — current data, memory, actions
    "what is the weather today",
    "what is the bitcoin price right now",
    "what are the latest news",
    "what did we talk about last time",
    "what is my name",
    "where do I live",
    "what do you know about me",
    "what did I tell you before",
    "what is the current temperature outside",
    "remind me what I said about my project",
    "search for the latest iPhone release",
    "find me a restaurant nearby",
    "turn on the lights in the living room",
    "turn off all lights",
    "set temperature to 22 degrees",
    "what is the current traffic to Moscow",
    "check if anyone is home",
    "what devices are currently on",
    "look up my public IP address",
    "show me recent news about Proxmox",
    # Russian — weather and commute
    "какая сегодня погода в Балашихе",
    "пойдет ли сегодня дождь",
    "какая температура на улице сейчас",
    "погода на завтра",
    "будет ли снег сегодня",
    "сколько ехать до Москвы сейчас",
    "какие пробки на дороге до Москвы",
    "время в пути на работу",
    "есть ли пробки сейчас",
    "стоит ли брать зонтик",
    # Russian — smart home control
    "включи свет в гостиной",
    "выключи свет на кухне",
    "какая температура дома",
    "установи температуру 22 градуса",
    "выключи все лампочки",
    "какие устройства сейчас включены",
    "включи ночной режим",
    "открой шторы в гостиной",
    "включи свет в спальне на 50 процентов",
    "выключи свет во всём доме",
    "включи вентилятор в детской",
    "закрыты ли все окна",
    "выключи телевизор",
    "какое потребление электричества сегодня",
    "включи кофемашину",
    "сколько у нас датчиков движения",
    "состояние всех дверных замков",
    "есть ли кто-нибудь дома",
    "установи будильник на 7 утра",
    # Russian — personal memory
    "как меня зовут",
    "где я живу",
    "что мы обсуждали в прошлый раз",
    "что ты знаешь о моем домашнем сервере",
    "напомни, какие сервисы я запускаю",
    "что я просил тебя запомнить",
    "что я говорил о своей сети",
    # Russian — current info lookups requiring network/tools
    "какой сейчас курс биткоина",
    "курс доллара к рублю сейчас",
    "какая последняя версия Docker",
    "как перезапустить Docker контейнер",
    "как посмотреть логи Docker контейнера",
    "какие новые функции в Home Assistant 2024",
    "есть ли проблемы у Cloudflare сегодня",
    "какие новые Zigbee устройства вышли в 2024 году",
    "найди хороший опенсорс менеджер фотографий",
    "последние новости Proxmox",
    "напиши bash команду для поиска больших файлов",
    "как вывести список всех запущенных контейнеров",
    "как проверить использование диска в Linux",
]

_COMPLEX_UTTERANCES = [
    # English
    "research everything about Elon Musk's recent projects and investments",
    "write a detailed report on climate change solutions with sources",
    "investigate the history and current state of quantum computing",
    "find and summarize the latest academic papers on transformer architectures",
    "analyze in depth the pros and cons of nuclear energy with citations",
    "research the background and controversies around this person",
    "compare all major cloud providers with detailed pricing and features",
    "write a comprehensive biography of this historical figure",
    "investigate what caused the 2008 financial crisis with multiple sources",
    "research the best programming languages in 2024 with detailed comparison",
    "find everything published about this medical condition and treatments",
    "do a deep dive into the latest developments in artificial general intelligence",
    "research and compare all options for starting a business in Europe",
    "investigate recent news and controversies around this company",
    "write a thorough analysis of geopolitical tensions in the Middle East",
    "find detailed information on the side effects and studies for this medication",
    "research the top 10 JavaScript frameworks with benchmarks and community data",
    "investigate who is funding AI research and what their goals are",
    "write a detailed market analysis for the electric vehicle industry",
    "research everything you can find about this startup or technology",
    # Russian — deep research
    "исследуй и сравни все варианты умного домашнего освещения",
    "напиши подробный отчет о протоколах умного дома",
    "изучи все самохостируемые медиасерверы и сравни их",
    "исследуй лучшие практики безопасности домашнего сервера",
    "сравни все системы резервного копирования для Linux",
    "напиши детальное сравнение WireGuard и OpenVPN",
    "исследуй все варианты голосового управления на русском языке",
    "изучи все опенсорс альтернативы Google сервисам",
    "напиши подробный анализ локальных языковых моделей",
    "исследуй лучшие инструменты мониторинга для домашнего сервера",
    # Russian — more deep research queries matching benchmark
    "исследуй и сравни Proxmox, Unraid и TrueNAS для домашней лаборатории",
    "напиши подробное руководство по безопасности домашнего сервера",
    "исследуй все доступные дашборды для самохостинга и сравни их",
    "найди детальные бенчмарки ARM одноплатных компьютеров для домашней лаборатории",
    "исследуй лучший стек мониторинга для самохостинга в 2024 году",
    "исследуй и сравни WireGuard, OpenVPN и Tailscale для домашней сети",
    "исследуй лучшие практики сегментации домашней сети с VLAN",
    "изучи все самохостируемые DNS решения и их возможности",
    "исследуй и сравни все платформы умного дома: Home Assistant и другие",
    "изучи лучшие Zigbee координаторы и их совместимость с Home Assistant",
    "напиши детальный отчет о поддержке протокола Matter и совместимости устройств",
    "исследуй все способы интеграции умных ламп с Home Assistant",
    "найди и сравни все варианты датчиков движения для умного дома",
    "исследуй и сравни все самохостируемые решения для хранения фотографий",
    "изучи лучшие самохостируемые медиасерверы: Jellyfin, Plex и Emby",
    "исследуй последние достижения в локальном LLM инференсе и обзор моделей",
    "изучи лучшие опенсорс альтернативы Google сервисов для приватности",
    "найди и кратко опиши все крупные самохостируемые менеджеры паролей",
    "напиши детальный анализ текущего состояния AI ассистентов для самохостинга",
    "исследуй и сравни все инструменты оркестрации контейнеров для домашней лаборатории",
    "изучи лучшие подходы к автоматическому резервному копированию в Linux",
    "исследуй и сравни все самохостируемые инструменты личных финансов",
    "изучи свежие CVE и уязвимости в популярном самохостируемом ПО",
    "напиши подробное руководство по настройке автоматизаций в Home Assistant",
    "исследуй все варианты голосового управления умным домом на русском языке",
    "сравни все системы резервного копирования для Linux: Restic, BorgBackup и другие",
    "исследуй лучшие самохостируемые системы мониторинга сети: Zabbix, Grafana",
    "изучи все варианты локального запуска языковых моделей на видеокарте",
    "напиши подробный отчет о технологиях синтеза речи с открытым исходным кодом",
    "исследуй все способы интеграции умных розеток с мониторингом потребления",
    "напиши полное руководство по настройке обратного прокси Caddy",
    "исследуй лучшие практики написания Docker Compose файлов для продакшена",
    "сравни все самохостируемые облачные хранилища: Nextcloud, Seafile и другие",
    "изучи все доступные локальные ассистенты с голосовым управлением",
    "исследуй все самохостируемые решения для блокировки рекламы: Pi-hole, AdGuard",
    "напиши детальное сравнение систем управления конфигурацией: Ansible, Puppet",
    "исследуй все протоколы умного дома и их плюсы и минусы: Zigbee, Z-Wave, Matter",
    "найди и сравни все фреймворки для создания локальных AI ассистентов",
    "исследуй лучшие решения для автоматического управления медиатекой",
    "изучи все варианты самохостируемых систем учёта расходов с возможностью импорта",
    "напиши сравнение всех вариантов самохостинга для хранения и синхронизации файлов",
    "исследуй все открытые протоколы для умного дома и их экосистемы",
    "изучи лучшие инструменты для автоматизации домашней инфраструктуры",
]

# Medium: queries that require tools, actions, or real-time data (not static knowledge)
_MEDIUM_PATTERNS = re.compile(
    r"(?:"
    # Russian smart home commands — always need HA integration
    r"(?:включи|выключи|открой|закрой|установи|поставь|убавь|прибавь|переключи)\s"
    r"|(?:какая|какой|какое|каково)\s+(?:температура|влажность|потребление|состояние|статус)\s"
    r"|(?:сколько|есть ли)\s.*(?:датчик|устройств|замк)"
    # Russian memory queries
    r"|как меня зовут|где я живу|что мы обсуждали|что я говорил|что я просил"
    r"|напомни\b|что ты знаешь обо мне"
    # Russian current info
    r"|курс (?:доллара|биткоина|евро|рубл)"
    r"|(?:последние |свежие )?новости\b"
    r"|(?:погода|температура)\s+(?:на завтра|на неделю)"
    # Smart home commands that don't use verb-first pattern
    r"|(?:свет|лампочк|освещени)\w*\s+(?:включ|выключ|убавь|прибавь)"
    r"|(?:дома|в доме|по всему дому)\s+(?:свет|лампочк)"
    r"|(?:режим|сцена)\s+(?:ночной|утренний|вечерний|кинотеатр)"
    r")",
    re.IGNORECASE,
)

LIGHT_REPLY_PROMPT = """You are a helpful Telegram assistant. Answer briefly and naturally (1-3 sentences). Be friendly."""

_EMBED_MODEL = "ollama/nomic-embed-text"


def _cosine(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


def _centroid(embeddings: list[list[float]]) -> list[float]:
    n = len(embeddings)
    dim = len(embeddings[0])
    return [sum(embeddings[i][d] for i in range(n)) / n for d in range(dim)]


def _format_history(history: list[dict]) -> str:
    if not history:
        return "(none)"
    lines = []
    for msg in history:
        role = msg.get("role", "?")
        content = str(msg.get("content", ""))[:200]
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


class Router:
    def __init__(
        self,
        model,
        embedder: AsyncOpenAI,
        fast_tool_runner: FastToolRunner | None = None,
    ):
        self.model = model  # qwen2.5:1.5b — used only for generating light replies
        self._embedder = embedder
        self._fast_tool_runner = fast_tool_runner
        self._light_centroid: list[float] | None = None
        self._medium_centroid: list[float] | None = None
        self._complex_centroid: list[float] | None = None

    async def initialize(self) -> None:
        """Pre-compute utterance embeddings. Call once at startup. Retries until LiteLLM is ready."""
        print("[router] embedding utterances for semantic classifier...", flush=True)
        texts = _LIGHT_UTTERANCES + _MEDIUM_UTTERANCES + _COMPLEX_UTTERANCES
        for attempt in range(10):
            try:
                resp = await self._embedder.embeddings.create(model=_EMBED_MODEL, input=texts)
                embeddings = [item.embedding for item in resp.data]
                n_light = len(_LIGHT_UTTERANCES)
                n_medium = len(_MEDIUM_UTTERANCES)
                self._light_centroid = _centroid(embeddings[:n_light])
                self._medium_centroid = _centroid(embeddings[n_light:n_light + n_medium])
                self._complex_centroid = _centroid(embeddings[n_light + n_medium:])
                print("[router] semantic classifier ready (3-tier)", flush=True)
                return
            except Exception as e:
                print(f"[router] embedding attempt {attempt+1}/10 failed: {e}", flush=True)
                await asyncio.sleep(3)
        print("[router] WARNING: could not initialize semantic classifier — will default to medium", flush=True)

    async def _classify_by_embedding(self, message: str) -> str:
        """Embed message and return 'light', 'medium', or 'complex' based on centroid similarity."""
        if self._light_centroid is None or self._medium_centroid is None or self._complex_centroid is None:
            return "medium"
        try:
            resp = await self._embedder.embeddings.create(model=_EMBED_MODEL, input=[message])
            emb = resp.data[0].embedding
            score_light = _cosine(emb, self._light_centroid)
            score_medium = _cosine(emb, self._medium_centroid)
            score_complex = _cosine(emb, self._complex_centroid)
            tier = max(
                [("light", score_light), ("medium", score_medium), ("complex", score_complex)],
                key=lambda x: x[1],
            )[0]
            print(
                f"[router] semantic: light={score_light:.3f} medium={score_medium:.3f} "
                f"complex={score_complex:.3f} → {tier}",
                flush=True,
            )
            return tier
        except Exception as e:
            print(f"[router] embedding classify error, defaulting to medium: {e}", flush=True)
            return "medium"

    async def route(
        self,
        message: str,
        history: list[dict],
    ) -> tuple[str, Optional[str]]:
        """
        Returns (tier, reply_or_None).
        For light tier: also generates the reply inline.
        For medium/complex: reply is None.
        """
        if self._fast_tool_runner and self._fast_tool_runner.any_matches(message.strip()):
            names = self._fast_tool_runner.matching_names(message.strip())
            print(f"[router] fast_tool_match={names} → medium", flush=True)
            return "medium", None

        if _LIGHT_PATTERNS.match(message.strip()):
            print("[router] regex→light", flush=True)
            return await self._generate_light_reply(message, history)

        if _COMPLEX_PATTERNS.search(message.strip()):
            print("[router] regex→complex", flush=True)
            return "complex", None

        if _MEDIUM_PATTERNS.search(message.strip()):
            print("[router] regex→medium", flush=True)
            return "medium", None

        tier = await self._classify_by_embedding(message)

        if tier != "light":
            return tier, None

        return await self._generate_light_reply(message, history)

    async def _generate_light_reply(
        self, message: str, history: list[dict]
    ) -> tuple[str, Optional[str]]:
        """Generate a short reply using qwen2.5:1.5b for light-tier messages."""
        history_text = _format_history(history)
        context = f"\nConversation history:\n{history_text}" if history else ""
        try:
            reply_response = await self.model.ainvoke([
                SystemMessage(content=LIGHT_REPLY_PROMPT + context),
                HumanMessage(content=message),
            ])
            reply_text = reply_response.content or ""
            reply_text = re.sub(r"<think>.*?</think>", "", reply_text, flags=re.DOTALL).strip()
            if not reply_text:
                print("[router] light reply empty, falling back to medium", flush=True)
                return "medium", None
            print(f"[router] light reply: {len(reply_text)} chars", flush=True)
            return "light", reply_text
        except Exception as e:
            print(f"[router] light reply error, falling back to medium: {e}", flush=True)
            return "medium", None