38 lines
1.4 KiB
Python
38 lines
1.4 KiB
Python
from app.services.smart_router import ENGINE_COST_PROFILES, SmartRouter
|
|
from app.services.api_key_manager import APIKeyManager
|
|
|
|
|
|
class EngineSelector:
|
|
def __init__(self, key_manager: APIKeyManager):
|
|
self.key_manager = key_manager
|
|
self.router = SmartRouter(key_manager=key_manager)
|
|
|
|
def select_engines(
|
|
self,
|
|
max_engines: int = 5,
|
|
prefer_domestic: bool = True,
|
|
min_cost_tier: str | None = None,
|
|
) -> list[str]:
|
|
engines = self.router.select_engines(max_engines, prefer_domestic)
|
|
|
|
if min_cost_tier:
|
|
tier_order = ["free", "low_cost", "mid_cost", "high_cost"]
|
|
if min_cost_tier in tier_order:
|
|
min_idx = tier_order.index(min_cost_tier)
|
|
filtered = []
|
|
for engine in engines:
|
|
profile = ENGINE_COST_PROFILES.get(engine)
|
|
if profile and tier_order.index(profile.cost_tier.value) >= min_idx:
|
|
filtered.append(engine)
|
|
return filtered
|
|
|
|
return engines
|
|
|
|
def get_best_value_engine(self) -> str | None:
|
|
available = self.router.get_available_engines()
|
|
for engine in available:
|
|
profile = ENGINE_COST_PROFILES.get(engine)
|
|
if profile and profile.cost_tier.value in ["free", "low_cost"]:
|
|
return engine
|
|
return available[0] if available else None
|