113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
|
|
class CostTier(str, Enum):
|
|
FREE = "free"
|
|
LOW_COST = "low_cost"
|
|
MID_COST = "mid_cost"
|
|
HIGH_COST = "high_cost"
|
|
|
|
|
|
@dataclass
|
|
class EngineCostProfile:
|
|
engine_type: str
|
|
cost_tier: CostTier
|
|
input_price_per_million: float
|
|
output_price_per_million: float
|
|
has_free_tier: bool
|
|
requires_own_key: bool
|
|
geo_relevance: int
|
|
domestic: bool
|
|
|
|
|
|
ENGINE_COST_PROFILES: dict[str, EngineCostProfile] = {
|
|
"deepseek": EngineCostProfile("deepseek", CostTier.FREE, 0.25, 6.0, True, False, 4, True),
|
|
"qwen": EngineCostProfile("qwen", CostTier.FREE, 0.3, 0.6, True, False, 4, True),
|
|
"wenxin": EngineCostProfile("wenxin", CostTier.FREE, 0.012, 0.012, True, False, 3, True),
|
|
"kimi": EngineCostProfile("kimi", CostTier.LOW_COST, 12.0, 12.0, True, False, 4, True),
|
|
"doubao": EngineCostProfile("doubao", CostTier.LOW_COST, 0.5, 0.9, True, False, 3, True),
|
|
"gemini": EngineCostProfile("gemini", CostTier.LOW_COST, 0.5, 2.0, True, False, 4, False),
|
|
"yuanbao": EngineCostProfile("yuanbao", CostTier.MID_COST, 0.8, 2.0, True, False, 3, True),
|
|
"chatgpt": EngineCostProfile("chatgpt", CostTier.HIGH_COST, 1.0, 4.0, False, True, 5, False),
|
|
"perplexity": EngineCostProfile("perplexity", CostTier.HIGH_COST, 35.0, 35.0, False, True, 5, False),
|
|
}
|
|
|
|
_TIER_ORDER = [CostTier.FREE, CostTier.LOW_COST, CostTier.MID_COST, CostTier.HIGH_COST]
|
|
|
|
|
|
def _get_profile(engine: str) -> EngineCostProfile | None:
|
|
return ENGINE_COST_PROFILES.get(engine)
|
|
|
|
|
|
class SmartRouter:
|
|
def __init__(self, available_engines: list[str] | None = None, user_engines: list[str] | None = None):
|
|
self.available_engines = available_engines or list(ENGINE_COST_PROFILES.keys())
|
|
self._user_engine_set = set(user_engines or [])
|
|
|
|
@property
|
|
def user_engines(self) -> list[str]:
|
|
return list(self._user_engine_set)
|
|
|
|
def select_engines(self, max_engines: int = 5, prefer_domestic: bool = True) -> list[str]:
|
|
tiers: dict[CostTier, list[str]] = {tier: [] for tier in CostTier}
|
|
for engine in self.available_engines:
|
|
profile = _get_profile(engine)
|
|
if profile is None:
|
|
continue
|
|
tiers[profile.cost_tier].append(engine)
|
|
|
|
for tier in _TIER_ORDER:
|
|
user = [e for e in tiers[tier] if e in self._user_engine_set]
|
|
non_user = [e for e in tiers[tier] if e not in self._user_engine_set]
|
|
tiers[tier] = user + non_user
|
|
|
|
if prefer_domestic:
|
|
for tier in _TIER_ORDER:
|
|
domestic = [e for e in tiers[tier] if _get_profile(e).domestic]
|
|
international = [e for e in tiers[tier] if not _get_profile(e).domestic]
|
|
tiers[tier] = domestic + international
|
|
|
|
selected: list[str] = []
|
|
for tier in _TIER_ORDER:
|
|
for engine in tiers[tier]:
|
|
if len(selected) >= max_engines:
|
|
return selected
|
|
if engine not in selected:
|
|
selected.append(engine)
|
|
|
|
return selected[:max_engines]
|
|
|
|
def get_cost_estimate(self, engines: list[str], estimated_input_tokens: int = 500, estimated_output_tokens: int = 1000) -> dict:
|
|
total_cost = 0.0
|
|
details: dict[str, dict] = {}
|
|
for engine in engines:
|
|
profile = _get_profile(engine)
|
|
if profile is None:
|
|
continue
|
|
cost = (estimated_input_tokens / 1_000_000 * profile.input_price_per_million
|
|
+ estimated_output_tokens / 1_000_000 * profile.output_price_per_million)
|
|
details[engine] = {"cost": round(cost, 6), "tier": profile.cost_tier.value}
|
|
total_cost += cost
|
|
return {"total_cost": round(total_cost, 6), "per_engine": details}
|
|
|
|
def _engines_by_tier(self, tier: CostTier) -> list[str]:
|
|
return [e for e in self.available_engines if (p := _get_profile(e)) is not None and p.cost_tier == tier]
|
|
|
|
def _engines_by_requires_key(self) -> list[str]:
|
|
return [e for e in self.available_engines if (p := _get_profile(e)) is not None and p.requires_own_key]
|
|
|
|
def get_recommended_combination(self) -> dict:
|
|
free_engines = self._engines_by_tier(CostTier.FREE)
|
|
low_cost = self._engines_by_tier(CostTier.LOW_COST)
|
|
user_only = self._engines_by_requires_key()
|
|
|
|
return {
|
|
"basic": free_engines[:3],
|
|
"standard": (free_engines + low_cost)[:5],
|
|
"premium": list(self.available_engines),
|
|
"user_premium": user_only,
|
|
}
|