Files
zk-data-agent/agent_platform/models.py
T
2026-07-26 06:52:12 +08:00

62 lines
1.5 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from typing import Literal
Mode = Literal["chat", "work"]
Strength = Literal["light", "medium", "high"]
@dataclass(frozen=True, slots=True)
class ModelSpec:
public_id: str
display_name: str
mode: Mode
strength: Strength
provider_model: str
max_iterations: int
context_char_budget: int
_TIERS = {
"light": ("轻度", "ChatGPT-5.6:Luna", 8, 120_000),
"medium": ("", "ChatGPT-5.6:Terra", 16, 240_000),
"high": ("", "ChatGPT-5.6:Sol", 24, 400_000),
}
MODEL_SPECS: dict[str, ModelSpec] = {
f"{mode}-{strength}": ModelSpec(
public_id=f"{mode}-{strength}",
display_name=f"{'Chat' if mode == 'chat' else 'Work'} · {label}",
mode=mode,
strength=strength,
provider_model=provider,
max_iterations=max_iterations,
context_char_budget=context_budget,
)
for mode in ("chat", "work")
for strength, (label, provider, max_iterations, context_budget) in _TIERS.items()
}
def get_model_spec(model_id: str) -> ModelSpec:
try:
return MODEL_SPECS[model_id]
except KeyError as exc:
raise ValueError(f"Unsupported model: {model_id}") from exc
def openai_model_list() -> dict:
return {
"object": "list",
"data": [
{
"id": spec.public_id,
"object": "model",
"owned_by": "k1412-agent",
"name": spec.display_name,
}
for spec in MODEL_SPECS.values()
],
}