unsloth/studio/backend/core/inference/providers.py
2026-05-12 11:42:10 +01:00

203 lines
7.2 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Static registry of supported external LLM providers.
All providers expose OpenAI-compatible /v1/chat/completions endpoints
with Bearer token authentication and SSE streaming support.
"""
import re
from typing import Any
PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"openai": {
"display_name": "OpenAI",
"base_url": "https://api.openai.com/v1",
"default_models": [
"gpt-5.5",
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.3",
"o3",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
# Keep the model picker scoped to the current generation. The remote
# /v1/models listing returns dozens of historical snapshots, fine-tunes
# and non-chat models (embeddings, TTS, image, moderation) that we
# never want to surface in the chat UI. Filtering here so backend
# is the single source of truth.
"model_id_allowlist": re.compile(r"^(gpt-5\.[345]|gpt-4\.5|o3)(?:[-.]|$)"),
# Hide dated snapshots
"model_id_denylist": re.compile(r"-\d{4}-\d{2}-\d{2}$"),
},
"anthropic": {
"display_name": "Anthropic",
"base_url": "https://api.anthropic.com/v1",
"default_models": [
"claude-opus-4-5",
"claude-sonnet-4-5",
"claude-haiku-4-5",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": False,
"auth_header": "x-api-key",
"auth_prefix": "",
"extra_headers": {
"anthropic-version": "2023-06-01",
},
"openai_compatible": False,
"notes": "Native Anthropic Messages API. Uses x-api-key header and /v1/messages endpoint with SSE translation.",
},
"gemini": {
"display_name": "Google Gemini",
"base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
"default_models": [
"gemini-3-flash-preview",
"gemini-2.5-flash",
"gemini-2.5-pro",
"gemini-2.5-flash-lite",
"gemini-2.0-flash",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.",
},
"deepseek": {
"display_name": "DeepSeek",
"base_url": "https://api.deepseek.com/v1",
"default_models": [
"deepseek-chat",
"deepseek-reasoner",
],
"supports_streaming": True,
"supports_vision": False,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "OpenAI-compatible API. deepseek-chat = V3, deepseek-reasoner = R1 thinking mode.",
},
"mistral": {
"display_name": "Mistral AI",
"base_url": "https://api.mistral.ai/v1",
"default_models": [
"mistral-large-2512",
"mistral-small-2506",
"codestral-latest",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
},
"kimi": {
"display_name": "Kimi",
"base_url": "https://api.moonshot.ai/v1",
"default_models": [
"kimi-k2-0711-preview",
"moonshot-v1-8k",
"moonshot-v1-32k",
"moonshot-v1-128k",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "Moonshot API key. China: use base URL https://api.moonshot.cn/v1",
},
"qwen": {
"display_name": "Qwen",
"base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
"default_models": [
"qwen-plus",
"qwen-turbo",
"qwen-max",
"qwen2.5-72b-instruct",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "DashScope API key. China mainland: override base URL to https://dashscope.aliyuncs.com/compatible-mode/v1",
},
"huggingface": {
"display_name": "Hugging Face",
"base_url": "https://router.huggingface.co/v1",
"default_models": [
"meta-llama/Llama-3.3-70B-Instruct",
"Qwen/Qwen2.5-72B-Instruct",
"deepseek-ai/DeepSeek-V3",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "HF token from huggingface.co/settings/tokens. Inference Providers OpenAI-compatible API.",
"model_list_mode": "curated",
},
"openrouter": {
"display_name": "OpenRouter",
"base_url": "https://openrouter.ai/api/v1",
"default_models": [
"openai/gpt-4o",
"google/gemini-2.5-flash",
"mistralai/mistral-small-3.1-24b-instruct",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"extra_headers": {
"HTTP-Referer": "https://unsloth.ai",
"X-Title": "Unsloth Studio",
},
"notes": "Unified gateway to 300+ models across all major providers. HTTP-Referer and X-Title headers sent for attribution.",
"model_list_mode": "curated",
},
}
def get_provider_info(provider_type: str) -> dict[str, Any] | None:
"""Return the registry entry for a provider type, or None if unknown."""
return PROVIDER_REGISTRY.get(provider_type)
def get_base_url(provider_type: str) -> str | None:
"""Return the default base URL for a provider type."""
info = PROVIDER_REGISTRY.get(provider_type)
return info["base_url"] if info else None
def list_available_providers() -> list[dict[str, Any]]:
"""Return all registered providers (for the /registry endpoint)."""
result = []
for provider_type, info in PROVIDER_REGISTRY.items():
result.append(
{
"provider_type": provider_type,
"display_name": info["display_name"],
"base_url": info["base_url"],
"default_models": info["default_models"],
"supports_streaming": info["supports_streaming"],
"supports_vision": info.get("supports_vision", False),
"supports_tool_calling": info.get("supports_tool_calling", False),
"model_list_mode": info.get("model_list_mode", "remote"),
}
)
return result