core/engine/src/llm/model_classifier.py - main - alexefan136/flowstack

core/engine/src/llm/model_classifier.py
318 строк10 KB
Alexander Efanov
upd fix
31 июл 2026, 19:17
31 июл 2026, 19:17d146d86
"""
Автоматический классификатор моделей по ID.

Определяет категорию, провайдера, теги, контекстное окно и возможности
модели по её ID с помощью pattern-matching.

Актуализирован под Cloud.ru Foundation Models:
  - 'Qwen/Qwen3-Embedding-0.6B'  → EMBEDDING
  - 'BAAI/bge-reranker-v2-m3'    → RERANKER (не embedding!)
  - 'deepseek-ai/DeepSeek-OCR-2' → OCR
  - 'openai/whisper-large-v3'    → STT
  - 'hivetrace/HiveTraceGuard-Pro' → GUARDRAILS
  - 'Qwen/Qwen3-Coder-Next'      → CODE
  - 'deepseek-ai/DeepSeek-R1-0528' → REASONING
"""

from __future__ import annotations

import re

from .model_registry import ModelCategory

# =============================================================================
# PROVIDER DETECTION
# =============================================================================

# Маппинг паттернов ID → провайдер (порядок важен: специфичные первыми)
PROVIDER_PATTERNS: list[tuple[str, str]] = [
    # DeepSeek
    ("deepseek", "deepseek"),
    # Qwen / Alibaba
    ("qwen", "qwen"),
    # OpenAI
    ("openai", "openai"),
    ("gpt", "openai"),
    ("chatgpt", "openai"),
    # Anthropic
    ("anthropic", "anthropic"),
    ("claude", "anthropic"),
    # Google
    ("google", "google"),
    ("gemini", "google"),
    # Meta
    ("meta-llama", "meta"),
    ("llama", "meta"),
    # Sber (GigaChat)
    ("ai-sage", "sber"),
    ("gigachat", "sber"),
    ("sber", "sber"),
    # MiniMax
    ("minimax", "minimax"),
    # Moonshot (Kimi)
    ("moonshot", "moonshot"),
    ("kimi", "moonshot"),
    # Zhipu AI (GLM)
    ("zai-org", "zhipu"),
    ("z-ai", "zhipu"),
    ("glm", "zhipu"),
    # BAAI (BGE)
    ("baai", "baai"),
    ("bge", "baai"),
    # HiveTrace (Guardrails)
    ("hivetrace", "hivetrace"),
    # Xiaomi
    ("xiaomi", "xiaomi"),
    ("mimo", "xiaomi"),
    # Meituan
    ("meituan", "meituan"),
    ("longcat", "meituan"),
]


def detect_provider(model_id: str) -> str:
    """
    Определить провайдера модели по ID.

    Сначала ищет по префиксу (до '/'), потом по всему ID.

    Examples:
        'deepseek-ai/DeepSeek-V4-Pro' → 'deepseek'
        'ai-sage/GigaChat3-10B-A1.8B' → 'sber'
        'BAAI/bge-m3'                 → 'baai'
    """
    lower = model_id.lower()
    prefix = lower.split("/")[0] if "/" in lower else lower

    # По префиксу (точнее)
    for pattern, provider in PROVIDER_PATTERNS:
        if pattern in prefix:
            return provider

    # Fallback: по всему ID
    for pattern, provider in PROVIDER_PATTERNS:
        if pattern in lower:
            return provider

    return "unknown"


# =============================================================================
# CLASSIFICATION RULES
# =============================================================================

# Правила классификации: (regex_pattern, category)
# Порядок ВАЖЕН — первое совпадение выигрывает.
# Критично: RERANKER до EMBEDDING (bge-reranker не должен стать embedding),
#           VL-варианты до общих категорий.
CLASSIFICATION_RULES: list[tuple[str, ModelCategory]] = [
    # OCR
    (r"(?i)ocr", ModelCategory.OCR),
    # Embeddings — VL-вариант до общего
    (r"(?i)vl[-_]?embedding|vl.*embedding", ModelCategory.EMBEDDING_VL),
    # Rerankers — VL-вариант до общего, ОБА до EMBEDDING (bge-reranker!)
    (r"(?i)vl[-_]?reranker|vl.*reranker", ModelCategory.RERANKER_VL),
    (r"(?i)reranker", ModelCategory.RERANKER),
    # Embeddings — общий (после reranker!)
    (r"(?i)embedding|\bbge\b|text[-_]?embedding", ModelCategory.EMBEDDING),
    # Speech
    (r"(?i)whisper|speech[-_]?to[-_]?text|\bstt\b|\basr\b", ModelCategory.STT),
    (r"(?i)\btts\b|text[-_]?to[-_]?speech|voice.*synth", ModelCategory.TTS),
    # Guardrails / Moderation (HiveTraceGuard содержит "guard")
    (r"(?i)moderation|guardrail|guard|hivetrace|content[-_]?filter", ModelCategory.GUARDRAILS),
    # Reasoning (chain-of-thought)
    (
        r"(?i)-R1\b|reasoning|thinking|\bthink\b|\bo3[-_]|\bo1[-_]|deep[-_]?research",
        ModelCategory.REASONING,
    ),
    # Code
    (r"(?i)codellama|coder|code-|codex|star.?coder", ModelCategory.CODE),
    # Vision (chat с vision, если явно указано)
    (r"(?i)-VL\b|vision", ModelCategory.CHAT),
]


def classify_model(model_id: str) -> ModelCategory:
    """
    Определить категорию модели по её ID.

    Args:
        model_id: ID модели (например, 'Qwen/Qwen3-Embedding-0.6B')

    Returns:
        ModelCategory. Если не удалось определить — CHAT по умолчанию.
    """
    for pattern, category in CLASSIFICATION_RULES:
        if re.search(pattern, model_id):
            return category
    return ModelCategory.CHAT


# =============================================================================
# TAG EXTRACTION
# =============================================================================

# Ключевые слова для тегов
_TAG_KEYWORDS = [
    "vl",
    "ocr",
    "embedding",
    "reranker",
    "whisper",
    "reasoning",
    "thinking",
    "coder",
    "codex",
    "flash",
    "guard",
    "research",
]


def extract_tags(model_id: str) -> list[str]:
    """
    Извлечь теги из ID модели.

    Examples:
        'Qwen/Qwen3-VL-Embedding-8B' → ['8b', 'vl', 'embedding', 'qwen', 'pro'?]
        'deepseek-ai/DeepSeek-V4-Pro' → ['v4', 'deepseek', 'pro']
    """
    tags: list[str] = []
    normalized = model_id.lower().replace("/", " ").replace("_", " ")

    # Размер модели (0.6B, 10B, 432B, 1.8B)
    size_match = re.search(r"(\d+(?:\.\d+)?)\s*[bB]\b", normalized)
    if size_match:
        tags.append(f"{size_match.group(1)}b")

    # Ключевые слова
    for kw in _TAG_KEYWORDS:
        if kw in normalized:
            tags.append(kw)

    # Провайдер
    provider = detect_provider(model_id)
    if provider != "unknown":
        tags.append(provider)

    # Специальные маркеры (варианты моделей)
    if "pro" in normalized:
        tags.append("pro")
    if "max" in normalized:
        tags.append("max")
    if any(marker in normalized for marker in ("mini", "lite", "nano", "flash")):
        tags.append("light")

    # Dedupe с сохранением порядка
    return list(dict.fromkeys(tags))


# =============================================================================
# CONTEXT WINDOW INFERENCE
# =============================================================================


def infer_context_window(model_id: str) -> int:
    """
    Оценить context window по ID модели.

    Использует явные паттерны, известные семейства и размер модели.
    Это fallback для авто-обнаруженных моделей (статический реестр
    имеет точные значения).
    """
    lower = model_id.lower()

    # Явные паттерны контекста
    if re.search(r"\b1m\b|1048576|1000000|1047576|1050000", lower):
        return 1_000_000
    if re.search(r"512k|524288", lower):
        return 524_288
    if re.search(r"256k|262144", lower):
        return 262_144
    if re.search(r"200k|200000|202752|204800", lower):
        return 200_000
    if re.search(r"128k|131072", lower):
        return 131_072
    if re.search(r"64k|65536", lower):
        return 65_536
    if re.search(r"32k|32768", lower):
        return 32_768
    if re.search(r"16k|16384", lower):
        return 16_384
    if re.search(r"\b8k\b|8192", lower):
        return 8_192

    # Известные семейства (эвристика)
    if any(
        p in lower
        for p in ("gpt-5", "gpt-4.1", "gemini-2.5", "gemini-3", "claude-opus", "claude-sonnet")
    ):
        return 200_000
    if "gpt-4o" in lower:
        return 128_000

    # По размеру модели
    size_match = re.search(r"(\d+(?:\.\d+)?)\s*[bB]\b", lower)
    if size_match:
        size_b = float(size_match.group(1))
        if size_b >= 70:
            return 131_072
        if size_b >= 30:
            return 65_536
        if size_b >= 7:
            return 32_768
        if size_b >= 1:
            return 8_192

    return 8_192


# =============================================================================
# CAPABILITIES INFERENCE
# =============================================================================

_GENERATIVE_CATEGORIES = (
    ModelCategory.CHAT,
    ModelCategory.REASONING,
    ModelCategory.CODE,
    ModelCategory.GUARDRAILS,
)

_VISION_CATEGORIES = (
    ModelCategory.OCR,
    ModelCategory.EMBEDDING_VL,
    ModelCategory.RERANKER_VL,
)


def infer_capabilities(model_id: str, category: ModelCategory) -> dict[str, bool]:
    """
    Определить возможности модели по ID и категории.

    Используется для авто-обнаруженных моделей (у которых нет явных
    метаданных из API).

    Returns:
        {
            "streaming": bool,
            "function_calling": bool,
            "vision": bool,
            "reasoning": bool,
            "supports_json": bool,
        }
    """
    lower = model_id.lower()
    is_generative = category in _GENERATIVE_CATEGORIES

    return {
        "streaming": is_generative,
        "function_calling": is_generative,
        "vision": category in _VISION_CATEGORIES or "vl" in lower or "vision" in lower,
        "reasoning": (
            category == ModelCategory.REASONING
            or any(
                marker in lower for marker in ("-r1", "thinking", "reasoning", "think", "o3", "o1")
            )
        ),
        "supports_json": is_generative,
    }