def get_router_llm(
primary: LLMConfig,
fallbacks: list[LLMConfig],
router_settings: dict | None = None,
) -> BaseChatModel:
"""Return a ``ChatLiteLLMRouter`` backed by a ``litellm.Router``.
Args:
primary: ``LLMConfig`` for the primary model.
fallbacks: Ordered list of ``LLMConfig`` fallback configs.
router_settings: Extra kwargs forwarded to ``litellm.Router``
(e.g. ``num_retries``).
"""
from langchain_litellm import ChatLiteLLMRouter # noqa: PLC0415
from datarobot_genai.core.router import build_litellm_router # noqa: PLC0415
class _CleanStreamRouter(ChatLiteLLMRouter):
"""Strip raw tool-call deltas from streaming ``additional_kwargs``.
``ChatLiteLLMRouter`` puts raw streaming tool-call delta objects into
``additional_kwargs["tool_calls"]``. When chunks accumulate these become
a flat list of partial deltas with fragmentary JSON arguments. The
correct data already lives in ``tool_call_chunks``; stripping the extra
key lets downstream code use that path instead.
Also normalizes list-form content (see ``_wrap_bare_text_blocks``) so the
router path gets the same protection as the non-router model.
"""
def _create_message_dicts(
self, messages: list[BaseMessage], stop: list[str] | None
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
message_dicts, params = super()._create_message_dicts(messages, stop)
return _wrap_bare_text_blocks(message_dicts), params
def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGenerationChunk]:
for chunk in super()._stream(*args, **kwargs):
chunk.message.additional_kwargs.pop("tool_calls", None)
yield chunk
async def _astream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ChatGenerationChunk]:
async for chunk in super()._astream(*args, **kwargs):
chunk.message.additional_kwargs.pop("tool_calls", None)
yield chunk
router = build_litellm_router(primary, fallbacks, router_settings)
return _CleanStreamRouter(router=router, model="primary", streaming=True)