from typing import Callable
class ModelRouter:
def __init__(self):
self.models = {}
self.routing_rules = []
def register_model(self, name: str, client: Callable, config: dict):
self.models[name] = {"client": client, "config": config}
def add_routing_rule(self, condition: Callable, model_name: str):
self.routing_rules.append((condition, model_name))
def route(self, request: dict) -> str:
"""Déterminer quel modèle utiliser."""
for condition, model_name in self.routing_rules:
if condition(request):
return model_name
return "default"
async def generate(self, request: dict) -> dict:
model_name = self.route(request)
model = self.models[model_name]
try:
response = await model["client"].generate(request["prompt"])
return {"response": response, "model": model_name}
except Exception as e:
return await self.fallback(request, exclude=[model_name])
router = ModelRouter()
router.register_model("gpt-4", gpt4_client, {"cost": "high", "quality": "best"})
router.register_model("gpt-3.5", gpt35_client, {"cost": "low", "quality": "good"})
router.register_model("claude", claude_client, {"cost": "medium", "quality": "best"})
router.add_routing_rule(
lambda r: len(r["prompt"]) < 100,
"gpt-3.5"
)