ForcePilot/backend/package/yuxi/channel/extensions/mattermost/model_picker.py

59 lines
1.9 KiB
Python
Raw Normal View History

from __future__ import annotations
import logging
logger = logging.getLogger(__name__)
DEFAULT_MODEL_OPTIONS = [
{"id": "gpt-4o", "name": "GPT-4o", "provider": "openai"},
{"id": "gpt-4-turbo", "name": "GPT-4 Turbo", "provider": "openai"},
{"id": "claude-3-opus", "name": "Claude 3 Opus", "provider": "anthropic"},
{"id": "claude-3-sonnet", "name": "Claude 3 Sonnet", "provider": "anthropic"},
{"id": "gemini-1.5-pro", "name": "Gemini 1.5 Pro", "provider": "google"},
]
class ModelPickerManager:
def get_available_models(self) -> list[dict]:
return DEFAULT_MODEL_OPTIONS
def get_model_info(self, model_id: str) -> dict | None:
for model in DEFAULT_MODEL_OPTIONS:
if model["id"] == model_id:
return model
return None
def build_picker_message(self, models: list[dict] | None = None) -> str:
if models is None:
models = DEFAULT_MODEL_OPTIONS
lines = ["**Available Models:**"]
for i, model in enumerate(models, 1):
lines.append(f"{i}. **{model['name']}** (`{model['id']}`) — {model['provider']}")
return "\n".join(lines)
def build_picker_buttons(
self,
models: list[dict] | None = None,
channel_id: str = "",
) -> list[dict]:
if models is None:
models = DEFAULT_MODEL_OPTIONS
buttons = []
for model in models:
buttons.append({
"id": f"select_model_{model['id']}",
"text": model["name"],
"name": model["name"],
"style": "default",
"context": {
"action": "select_model",
"model_id": model["id"],
"provider": model["provider"],
"__openclaw_channel_id": channel_id,
},
})
return buttons