builder / models.py
mgbam's picture
Update models.py
b412ca5 verified
raw
history blame
3.73 kB
# models.py
from dataclasses import dataclass
from typing import List, Optional
@dataclass
class ModelInfo:
"""
Represents metadata for an inference model.
Attributes:
name: Human‑readable name of the model.
id: Unique model identifier (HF or routed provider).
description: Short description of the model's capabilities.
default_provider: Preferred inference provider
("auto", "groq", "openai", "gemini", "fireworks").
"""
name: str
id: str
description: str
default_provider: str = "auto"
# Registry of supported models
AVAILABLE_MODELS: List[ModelInfo] = [
ModelInfo(
name="Moonshot Kimi‑K2",
id="moonshotai/Kimi‑K2‑Instruct",
description="Moonshot AI Kimi‑K2‑Instruct model for code generation and general tasks",
default_provider="groq"
),
ModelInfo(
name="DeepSeek V3",
id="deepseek‑ai/DeepSeek‑V3‑0324",
description="DeepSeek V3 model for code generation"
),
ModelInfo(
name="DeepSeek R1",
id="deepseek‑ai/DeepSeek‑R1‑0528",
description="DeepSeek R1 model for code generation"
),
ModelInfo(
name="ERNIE‑4.5‑VL",
id="baidu/ERNIE‑4.5‑VL‑424B‑A47B‑Base‑PT",
description="ERNIE‑4.5‑VL model for multimodal code generation with image support"
),
ModelInfo(
name="MiniMax M1",
id="MiniMaxAI/MiniMax‑M1‑80k",
description="MiniMax M1 model for code generation and general tasks"
),
ModelInfo(
name="Qwen3‑235B‑A22B",
id="Qwen/Qwen3‑235B‑A22B",
description="Qwen3‑235B‑A22B model for code generation and general tasks"
),
ModelInfo(
name="Qwen3‑235B‑A22B‑Instruct‑2507",
id="Qwen/Qwen3‑235B‑A22B‑Instruct‑2507",
description="Qwen3‑235B‑A22B‑Instruct‑2507 model for code generation and general tasks"
),
ModelInfo(
name="Qwen3‑Coder‑480B‑A35B",
id="Qwen/Qwen3‑Coder‑480B‑A35B‑Instruct",
description="Qwen3‑Coder‑480B‑A35B‑Instruct model for advanced code generation"
),
ModelInfo(
name="Qwen3‑32B",
id="Qwen/Qwen3‑32B",
description="Qwen3‑32B model for large‑scale code and general tasks"
),
ModelInfo(
name="SmolLM3‑3B",
id="HuggingFaceTB/SmolLM3‑3B",
description="SmolLM3‑3B model for code generation and general tasks"
),
ModelInfo(
name="GLM‑4.1V‑9B‑Thinking",
id="THUDM/GLM‑4.1V‑9B‑Thinking",
description="GLM‑4.1V‑9B‑Thinking model for multimodal code generation with image support"
),
ModelInfo(
name="OpenAI GPT‑4",
id="openai/gpt‑4",
description="OpenAI GPT‑4 via HF Inference Providers",
default_provider="openai"
),
ModelInfo(
name="Gemini Pro",
id="gemini/pro",
description="Google Gemini Pro via HF Inference Providers",
default_provider="gemini"
),
ModelInfo(
name="Fireworks AI",
id="fireworks‑ai/fireworks‑v1",
description="Fireworks AI via HF Inference Providers",
default_provider="fireworks"
),
]
def find_model(identifier: str) -> Optional[ModelInfo]:
"""
Lookup a model by its human name (case‑insensitive) or identifier.
Returns:
The matching ModelInfo or None if not found.
"""
identifier_lower = identifier.lower()
for m in AVAILABLE_MODELS:
if m.id == identifier or m.name.lower() == identifier_lower:
return m
return None