datarobot_genai.langgraph.llm

llm

get_router_llm

get_router_llm(primary: LLMConfig, fallbacks: list[LLMConfig], router_settings: dict | None = None) -> BaseChatModel

Return a ChatLiteLLMRouter backed by a litellm.Router.

Parameters:

Name	Type	Description	Default
primary	LLMConfig	LLMConfig for the primary model.	required
fallbacks	list[LLMConfig]	Ordered list of LLMConfig fallback configs.	required
router_settings	dict | None	Extra kwargs forwarded to litellm.Router (e.g. num_retries).	None

Source code in datarobot_genai/langgraph/llm.py

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)

get_llm

get_llm(model_name: str | None = None, parameters: dict | None = None, streaming: bool = True, reasoning: bool = False, config: LLMConfig | None = None) -> BaseChatModel

Build the framework LLM for a given (or the default) LLM config.

Pass config (for example resolve_config().resolve_llm_config(name="bob")) to build a specific configured LLM instance; omit it to build the app's default LLM.

Source code in datarobot_genai/langgraph/llm.py

def get_llm(
    model_name: str | None = None,
    parameters: dict | None = None,
    streaming: bool = True,
    reasoning: bool = False,
    config: LLMConfig | None = None,
) -> BaseChatModel:
    """Build the framework LLM for a given (or the default) LLM config.

    Pass ``config`` (for example ``resolve_config().resolve_llm_config(name="bob")``)
    to build a specific configured LLM instance; omit it to build the app's default LLM.
    """
    if config is None:
        config = resolve_config().resolve_llm_config(name=registered_default_llm_name())
    llm_type = config.get_llm_type()
    if llm_type == LLMType.GATEWAY:
        return get_datarobot_gateway_llm(model_name, parameters, streaming, reasoning)
    elif llm_type == LLMType.DEPLOYMENT:
        return get_datarobot_deployment_llm(
            config.llm_deployment_id,  # type: ignore[arg-type]
            model_name,
            parameters,
            streaming,
            reasoning,
        )
    elif llm_type == LLMType.NIM:
        return get_datarobot_nim_llm(
            config.llm_nim_deployment_id,  # type: ignore[arg-type]
            model_name,
            parameters,
            streaming,
            reasoning,
        )
    elif llm_type == LLMType.EXTERNAL:
        return get_external_llm(model_name, parameters, streaming, reasoning)
    else:
        raise ValueError(f"Invalid LLM type inferred from config: {llm_type}, config: {config}")