scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

working auto router

Browse source
This commit is contained in:
Ishaan Jaff 2025-07-24 15:30:35 -07:00
parent 334b90aad8
commit 1904e36e13
2 changed files with 17 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/router.py
View file

@ -165,12 +165,16 @@ from .router_utils.pattern_match_deployments import PatternMatchRouter
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.router_strategy.auto_router.auto_router import AutoRouter
from litellm.router_strategy.auto_router.auto_router import (
AutoRouter,
PreRoutingHookResponse,
)
Span = Union[_Span, Any]
else:
Span = Any
AutoRouter = Any
PreRoutingHookResponse = Any
class RoutingArgs(enum.Enum):
@ -6517,14 +6521,20 @@ class Router:
#########################################################
# Execute Pre-Routing Hooks
# this hook can modify the model, messages before the routing decision is made
#########################################################
await self.async_pre_routing_hook(
pre_routing_hook_response = await self.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
if pre_routing_hook_response is not None:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
#########################################################
healthy_deployments = await self.async_get_healthy_deployments(
@ -6644,7 +6654,7 @@ class Router:
messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
):
) -> Optional[PreRoutingHookResponse]:
"""
This hook is called before the routing decision is made.
@ -6654,7 +6664,7 @@ class Router:
# Check if any auto-router should be used
#########################################################
if model in self.auto_routers:
await self.auto_routers[model].async_pre_routing_hook(
return await self.auto_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
@ -6662,6 +6672,7 @@ class Router:
specific_deployment=specific_deployment,
)
return None

2
litellm/router_strategy/auto_router/auto_router.py
View file

@ -3,6 +3,7 @@ Auto-Routing Strategy that works with a Semantic Router Config
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
if TYPE_CHECKING:
@ -88,6 +89,7 @@ class AutoRouter(CustomLogger):
user_message: Dict[str, str] = messages[-1]
message_content: str = user_message.get("content", "")
route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer(text=message_content)
verbose_router_logger.debug(f"route_choice: {route_choice}")
if isinstance(route_choice, RouteChoice):
model = route_choice.name or self.default_model
elif isinstance(route_choice, list):