ForcePilot/backend/package/yuxi/channels/adapters/mattermost/model_picker.py

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from enum import StrEnum

MODEL_PICKER_PAGE_SIZE = 8


class PickerState(StrEnum):
    PROVIDERS = "providers"
    LIST = "list"
    SELECT = "select"
    BACK = "back"


@dataclass
class ModelOption:
    model_id: str
    display_name: str
    provider: str = ""
    description: str = ""

    @property
    def button_id(self) -> str:
        return hashlib.sha256(self.model_id.encode()).hexdigest()[:8]


def get_default_model_options() -> list[ModelOption]:
    return [
        ModelOption("gpt-4o-mini", "GPT-4o Mini", "openai", "轻量快速模型"),
        ModelOption("gpt-4o", "GPT-4o", "openai", "旗舰多模态模型"),
        ModelOption("claude-sonnet-4-20250514", "Claude Sonnet 4", "anthropic", "高性能推理模型"),
        ModelOption("deepseek-v3", "DeepSeek V3", "deepseek", "国产高性能模型"),
    ]


def get_available_providers(models: list[ModelOption] | None = None) -> list[str]:
    if models is None:
        models = get_default_model_options()
    seen: dict[str, bool] = {}
    result: list[str] = []
    for m in models:
        if m.provider and m.provider not in seen:
            seen[m.provider] = True
            result.append(m.provider)
    return result


def build_model_picker_actions(
    models: list[ModelOption] | None = None,
    callback_id: str = "model_picker",
    state: PickerState = PickerState.LIST,
    page: int = 0,
) -> list[dict]:
    if models is None:
        models = get_default_model_options()

    actions: list[dict] = []

    if state == PickerState.PROVIDERS:
        providers = get_available_providers(models)
        for provider in providers:
            actions.append(
                {
                    "id": f"provider_{provider}",
                    "name": f"查看 {provider} 模型",
                    "integration": {
                        "url": "",
                        "context": {
                            "action": "list_provider",
                            "provider": provider,
                            "callback_id": callback_id,
                            "state": PickerState.LIST.value,
                        },
                    },
                    "type": "button",
                    "text": f"🤖 {provider.upper()}",
                    "style": "",
                }
            )
        actions.append(
            {
                "id": "view_all",
                "name": "查看全部模型",
                "integration": {
                    "url": "",
                    "context": {
                        "action": "list_all",
                        "callback_id": callback_id,
                        "state": PickerState.LIST.value,
                    },
                },
                "type": "button",
                "text": "📋 全部模型",
                "style": "",
            }
        )
        return actions

    start = page * MODEL_PICKER_PAGE_SIZE
    page_models = models[start : start + MODEL_PICKER_PAGE_SIZE]
    total_pages = (len(models) + MODEL_PICKER_PAGE_SIZE - 1) // MODEL_PICKER_PAGE_SIZE

    for m in page_models:
        actions.append(
            {
                "id": f"model_{m.button_id}",
                "name": f"选择 {m.display_name}",
                "integration": {
                    "url": "",
                    "context": {
                        "action": "select_model",
                        "model_id": m.model_id,
                        "callback_id": callback_id,
                        "state": PickerState.SELECT.value,
                    },
                },
                "type": "button",
                "text": f"🤖 {m.display_name}",
                "style": "",
            }
        )

    if total_pages > 1:
        nav_actions: list[dict] = []
        if page > 0:
            nav_actions.append(
                {
                    "id": "prev_page",
                    "name": "上一页",
                    "integration": {
                        "url": "",
                        "context": {
                            "action": "page",
                            "page": str(page - 1),
                            "callback_id": callback_id,
                            "state": PickerState.LIST.value,
                        },
                    },
                    "type": "button",
                    "text": "⬅️ 上一页",
                    "style": "",
                }
            )
        if page < total_pages - 1:
            nav_actions.append(
                {
                    "id": "next_page",
                    "name": "下一页",
                    "integration": {
                        "url": "",
                        "context": {
                            "action": "page",
                            "page": str(page + 1),
                            "callback_id": callback_id,
                            "state": PickerState.LIST.value,
                        },
                    },
                    "type": "button",
                    "text": "➡️ 下一页",
                    "style": "",
                }
            )
        actions.extend(nav_actions)

    actions.append(
        {
            "id": "back_to_providers",
            "name": "按提供商筛选",
            "integration": {
                "url": "",
                "context": {
                    "action": "back",
                    "callback_id": callback_id,
                    "state": PickerState.PROVIDERS.value,
                },
            },
            "type": "button",
            "text": "🔙 按提供商筛选",
            "style": "",
        }
    )

    return actions


def build_model_picker_attachment(
    current_model: str = "",
    callback_id: str = "model_picker",
    state: PickerState = PickerState.PROVIDERS,
    page: int = 0,
    provider_filter: str = "",
) -> dict:
    models = get_default_model_options()

    if provider_filter:
        models = [m for m in models if m.provider == provider_filter]

    if state == PickerState.PROVIDERS:
        providers = get_available_providers(models)
        lines = [f"• **{p.upper()}**" for p in providers]
        text = "按 AI 提供商筛选模型：\n\n" + "\n".join(lines)
    else:
        start = page * MODEL_PICKER_PAGE_SIZE
        page_models = models[start : start + MODEL_PICKER_PAGE_SIZE]
        total_pages = (len(models) + MODEL_PICKER_PAGE_SIZE - 1) // MODEL_PICKER_PAGE_SIZE

        lines = []
        for m in page_models:
            marker = " ✅" if m.model_id == current_model else ""
            lines.append(f"• **{m.display_name}** ({m.provider}) — {m.description}{marker}")

        text = f"当前可用模型 (页 {page + 1}/{total_pages}):\n\n" + "\n".join(lines)

    actions = build_model_picker_actions(models, callback_id, state, page)

    return {
        "fallback": "模型选择器",
        "title": "🤖 模型选择",
        "text": text,
        "callback_id": callback_id,
        "actions": actions,
    }


def build_model_select_attachment(provider: str, callback_id: str = "model_select") -> dict:
    all_models = get_default_model_options()
    provider_models = [m for m in all_models if m.provider == provider]
    actions = build_model_picker_actions(provider_models, callback_id, PickerState.LIST, 0)

    return {
        "fallback": f"选择 {provider} 模型",
        "title": f"🤖 {provider} 模型",
        "text": f"选择 {provider} 下的模型：",
        "callback_id": callback_id,
        "actions": actions,
    }