OpenMesh/coworker/providers/capabilities.py

"""Per-model capability probe.

A heuristic table for now (refined as we probe real providers/endpoints). Accepts
either bare model names (`gpt-5.5`) or provider-qualified ones (`openai:gpt-5.5`).

Custom user-added models can have their capabilities overridden via
`set_custom_capabilities()`, which is populated from preferences.
"""

from __future__ import annotations

import dataclasses
from typing import Optional

from .base import ModelCapabilities

# Per-model custom capability overrides, keyed by full model id (e.g. "openai:smesh-smartops").
# Set by the runtime from user preferences so custom models can declare vision/pdf support.
_custom_caps: dict[str, dict[str, bool]] = {}


def set_custom_capabilities(model: str, caps: dict[str, bool]) -> None:
    """Register or update custom capability flags for a user-added model."""
    _custom_caps[model] = caps


def get_custom_capabilities(model: str) -> Optional[dict[str, bool]]:
    return _custom_caps.get(model)


def clear_custom_capabilities() -> None:
    _custom_caps.clear()


def _apply_custom_overrides(model: str, caps: ModelCapabilities) -> ModelCapabilities:
    """Apply user-configured capability overrides on top of heuristically-detected ones.

    ModelCapabilities is a frozen dataclass, so we use dataclasses.replace()
    to create a new instance with updated fields rather than mutating in place.
    """
    custom = _custom_caps.get(model)
    if not custom:
        return caps
    updates: dict[str, bool] = {}
    for key in ("vision", "pdf", "tools", "parallel_tool_calls", "streaming"):
        if key in custom:
            updates[key] = bool(custom[key])
    if not updates:
        return caps
    return dataclasses.replace(caps, **updates)


def capabilities_for(model: str) -> ModelCapabilities:
    # Curated models answer from the matrix (exact full-id match — including reseller ids
    # like `together:zai-org/GLM-5.2`, whose names defeat the prefix heuristics below).
    # Custom user-added models fall through to the heuristics, at their own risk.
    from .matrix import entry_for

    entry = entry_for(model)
    if entry is not None:
        return _apply_custom_overrides(model, entry.caps)

    provider = model.split(":", 1)[0].lower() if ":" in model else ""
    name = model.split(":", 1)[-1].lower()  # strip a provider prefix if present

    # Ollama (local) models vary widely and many fake/mishandle parallel tool calls — assume
    # tools work (we only point at tool-capable models) but stay conservative otherwise.
    # Vision is detected from common model naming conventions (-vl, vision, llava, etc.).
    if provider == "ollama":
        _vision_patterns = ("-vl", "vision", "llava", "bakllava", "cogvlm", "minicpm-v")
        has_vision = any(p in name for p in _vision_patterns)
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=has_vision, parallel_tool_calls=False, streaming=True
            ),
        )

    # Cloud-account providers (custom-added ids; curated ones answered from the matrix).
    # The family segment decides: Claude keeps its native capabilities; everything else
    # stays conservative until probed (Converse tool calling works across families, but
    # parallel calls and vision vary per model).
    if provider in ("bedrock", "vertex"):
        if name.startswith(("claude/", "gemini/")):
            return _apply_custom_overrides(
                model,
                ModelCapabilities(
                    tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
                ),
            )
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=False, parallel_tool_calls=False, streaming=True
            ),
        )

    # Claude / Gemini (both native): tools + vision + parallel tool calls + streaming. The
    # engine executes parallel calls sequentially and each converter folds the results into
    # the single next user message — exactly what both APIs require.
    if provider in ("anthropic", "gemini"):
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
            ),
        )

    # Modern OpenAI GPT models: tools + vision + parallel tool calls + streaming.
    if name.startswith(("gpt-5", "gpt-4")):
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
            ),
        )

    # OpenAI reasoning models: tools yes, parallel tool calls constrained.
    if name.startswith(("o1", "o3", "o4")):
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=False, parallel_tool_calls=False, streaming=True
            ),
        )

    # OpenAI-compatible vendors (DeepSeek, Z AI/GLM, Kimi, MiniMax, Qwen, xAI/Grok, Mistral):
    # tool calling + streaming across their current lineups; vision left off until probed
    # per-model (several have vision variants, but the text flagships are what we suggest).
    # Custom overrides can flip vision on for user-added fine-tunes like smesh-smartops.
    if name.startswith(
        ("deepseek", "glm", "kimi", "minimax", "qwen", "grok", "mistral", "magistral")
    ):
        return _apply_custom_overrides(
            model,
            ModelCapabilities(
                tools=True, vision=False, parallel_tool_calls=True, streaming=True
            ),
        )

    # Conservative default for unknown models.
    return _apply_custom_overrides(
        model,
        ModelCapabilities(
            tools=True, vision=False, parallel_tool_calls=False, streaming=True
        ),
    )