datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/nat/datarobot_auth_provider.py

@@ -0,0 +1,110 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
+from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from nat.authentication.interfaces import AuthProviderBase
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import HeaderCred
+from pydantic import Field
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+
+class Config(DataRobotAppFrameworkBaseSettings):
+    """
+    Finds variables in the priority order of: env
+    variables (including Runtime Parameters), .env, file_secrets, then
+    Pulumi output variables.
+    """
+
+    datarobot_api_token: str | None = None
+
+
+config = Config()
+
+
+class DataRobotAPIKeyAuthProviderConfig(APIKeyAuthProviderConfig, name="datarobot_api_key"):  # type: ignore[call-arg]
+    raw_key: str = Field(
+        description=(
+            "Raw API token or credential to be injected into the request parameter. "
+            "Used for 'bearer','x-api-key','custom', and other schemes. "
+        ),
+        default=config.datarobot_api_token,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+@register_auth_provider(config_type=DataRobotAPIKeyAuthProviderConfig)
+async def datarobot_api_key_client(
+    config: DataRobotAPIKeyAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[APIKeyAuthProvider]:
+    yield APIKeyAuthProvider(config=config)
+
+
+mcp_config = MCPConfig().server_config
+
+
+class DataRobotMCPAuthProviderConfig(AuthProviderBaseConfig, name="datarobot_mcp_auth"):  # type: ignore[call-arg]
+    headers: dict[str, str] | None = Field(
+        description=("Headers to be used for authentication. "),
+        default=mcp_config["headers"] if mcp_config else None,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+class DataRobotMCPAuthProvider(AuthProviderBase[DataRobotMCPAuthProviderConfig]):
+    def __init__(
+        self, config: DataRobotMCPAuthProviderConfig, config_name: str | None = None
+    ) -> None:
+        assert isinstance(config, DataRobotMCPAuthProviderConfig), (
+            "Config is not DataRobotMCPAuthProviderConfig"
+        )
+        super().__init__(config)
+
+    async def authenticate(self, user_id: str | None = None, **kwargs: Any) -> AuthResult | None:
+        """
+        Authenticate the user using the API key credentials.
+
+        Args:
+            user_id (str): The user ID to authenticate.
+
+        Returns
+        -------
+            AuthenticatedContext: The authenticated context containing headers
+        """
+        return AuthResult(
+            credentials=[
+                HeaderCred(name=name, value=value) for name, value in self.config.headers.items()
+            ]
+        )
+
+
+@register_auth_provider(config_type=DataRobotMCPAuthProviderConfig)
+async def datarobot_mcp_auth_provider(
+    config: DataRobotMCPAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[DataRobotMCPAuthProvider]:
+    yield DataRobotMCPAuthProvider(config=config)

datarobot_genai/nat/datarobot_llm_clients.py

@@ -0,0 +1,318 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import TypeVar
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_llm_client
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+from nat.utils.exception_handlers.automatic_retries import patch_with_retry
+
+from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
+from ..nat.datarobot_llm_providers import DataRobotLLMDeploymentModelConfig
+from ..nat.datarobot_llm_providers import DataRobotLLMGatewayModelConfig
+from ..nat.datarobot_llm_providers import DataRobotNIMModelConfig
+
+if TYPE_CHECKING:
+    from crewai import LLM
+    from langchain_openai import ChatOpenAI
+    from llama_index.llms.litellm import LiteLLM
+
+ModelType = TypeVar("ModelType")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    if isinstance(llm_config, RetryMixin):
+        client = patch_with_retry(
+            client,
+            retries=llm_config.num_retries,
+            retry_codes=llm_config.retry_on_status_codes,
+            retry_on_messages=llm_config.retry_on_errors,
+        )
+
+    return client
+
+
+def _create_datarobot_chat_openai(config: dict[str, Any]) -> Any:
+    from langchain_openai import ChatOpenAI  # noqa: PLC0415
+
+    class DataRobotChatOpenAI(ChatOpenAI):
+        def _get_request_payload(  # type: ignore[override]
+            self,
+            *args: Any,
+            **kwargs: Any,
+        ) -> dict:
+            # We need to default to include_usage=True for streaming but we get 400 response
+            # if stream_options is present for a non-streaming call.
+            payload = super()._get_request_payload(*args, **kwargs)
+            if not payload.get("stream"):
+                payload.pop("stream_options", None)
+            return payload
+
+    return DataRobotChatOpenAI(**config)
+
+
+def _create_datarobot_litellm(config: dict[str, Any]) -> Any:
+    from llama_index.core.base.llms.types import LLMMetadata  # noqa: PLC0415
+    from llama_index.llms.litellm import LiteLLM  # noqa: PLC0415
+
+    class DataRobotLiteLLM(LiteLLM):  # type: ignore[misc]
+        """DataRobotLiteLLM is a small LiteLLM wrapper class that makes all LiteLLM endpoints
+        compatible with the LlamaIndex library.
+        """
+
+        @property
+        def metadata(self) -> LLMMetadata:
+            """Returns the metadata for the LLM.
+
+            This is required to enable the is_chat_model and is_function_calling_model, which are
+            mandatory for LlamaIndex agents. By default, LlamaIndex assumes these are false unless
+            each individual model config in LiteLLM explicitly sets them to true. To use custom LLM
+            endpoints with LlamaIndex agents, you must override this method to return the
+            appropriate metadata.
+            """
+            return LLMMetadata(
+                context_window=128000,
+                num_output=self.max_tokens or -1,
+                is_chat_model=True,
+                is_function_calling_model=True,
+                model_name=self.model,
+            )
+
+    return DataRobotLiteLLM(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMGatewayModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN
+)
+async def datarobot_llm_gateway_langchain(
+    llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    config["base_url"] = config["base_url"] + "/genai/llmgw"
+    config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
+    client = _create_datarobot_chat_openai(config)
+    yield langchain_patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMGatewayModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI
+)
+async def datarobot_llm_gateway_crewai(
+    llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["base_url"] = config["base_url"].removesuffix("/api/v2")
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMGatewayModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX
+)
+async def datarobot_llm_gateway_llamaindex(
+    llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
+) -> AsyncGenerator[LiteLLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
+    client = _create_datarobot_litellm(config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMDeploymentModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN
+)
+async def datarobot_llm_deployment_langchain(
+    llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
+    config = llm_config.model_dump(
+        exclude={"type", "thinking"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
+    client = _create_datarobot_chat_openai(config)
+    yield langchain_patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMDeploymentModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI
+)
+async def datarobot_llm_deployment_crewai(
+    llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
+    config = llm_config.model_dump(
+        exclude={"type", "thinking"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["api_base"] = config.pop("base_url") + "/chat/completions"
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMDeploymentModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX
+)
+async def datarobot_llm_deployment_llamaindex(
+    llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
+) -> AsyncGenerator[LiteLLM]:
+    config = llm_config.model_dump(
+        exclude={"type", "thinking"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["api_base"] = config.pop("base_url") + "/chat/completions"
+    client = _create_datarobot_litellm(config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def datarobot_nim_langchain(
+    llm_config: DataRobotNIMModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
+    config = llm_config.model_dump(
+        exclude={"type", "thinking"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
+    client = _create_datarobot_chat_openai(config)
+    yield langchain_patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def datarobot_nim_crewai(
+    llm_config: DataRobotNIMModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
+    config = llm_config.model_dump(
+        exclude={"type", "thinking", "max_retries"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["api_base"] = config.pop("base_url") + "/chat/completions"
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def datarobot_nim_llamaindex(
+    llm_config: DataRobotNIMModelConfig, builder: Builder
+) -> AsyncGenerator[LiteLLM]:
+    config = llm_config.model_dump(
+        exclude={"type", "thinking"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    config["api_base"] = config.pop("base_url") + "/chat/completions"
+    client = _create_datarobot_litellm(config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN
+)
+async def datarobot_llm_component_langchain(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    from nat.plugins.langchain.llm import (  # noqa: PLC0415
+        _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+    )
+
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if config["use_datarobot_llm_gateway"]:
+        config["base_url"] = config["base_url"] + "/genai/llmgw"
+    config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
+    config.pop("use_datarobot_llm_gateway")
+    client = _create_datarobot_chat_openai(config)
+    yield langchain_patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI
+)
+async def datarobot_llm_component_crewai(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    from crewai import LLM  # noqa: PLC0415
+
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["base_url"] = config["base_url"].removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX
+)
+async def datarobot_llm_component_llamaindex(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LiteLLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    client = _create_datarobot_litellm(config)
+    yield _patch_llm_based_on_config(client, config)

datarobot_genai/nat/datarobot_llm_providers.py

@@ -0,0 +1,130 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.llm.openai_llm import OpenAIModelConfig
+from pydantic import AliasChoices
+from pydantic import Field
+
+
+class Config(DataRobotAppFrameworkBaseSettings):
+    """
+    Finds variables in the priority order of: env
+    variables (including Runtime Parameters), .env, file_secrets, then
+    Pulumi output variables.
+    """
+
+    datarobot_endpoint: str = "https://app.datarobot.com/api/v2"
+    datarobot_api_token: str | None = None
+    llm_deployment_id: str | None = None
+    nim_deployment_id: str | None = None
+    use_datarobot_llm_gateway: bool = False
+    llm_default_model: str | None = None
+
+
+config = Config()
+
+
+class DataRobotLLMComponentModelConfig(OpenAIModelConfig, name="datarobot-llm-component"):  # type: ignore[call-arg]
+    """A DataRobot LLM provider to be used with an LLM client."""
+
+    api_key: str | None = Field(
+        default=config.datarobot_api_token, description="DataRobot API key."
+    )
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint.rstrip("/")
+        if config.use_datarobot_llm_gateway
+        else config.datarobot_endpoint + f"/deployments/{config.llm_deployment_id}",
+        description="DataRobot LLM URL.",
+    )
+    model_name: str = Field(
+        validation_alias=AliasChoices("model_name", "model"),
+        serialization_alias="model",
+        description="The model name.",
+        default=config.llm_default_model or "datarobot-deployed-llm",
+    )
+    use_datarobot_llm_gateway: bool = config.use_datarobot_llm_gateway
+
+
+@register_llm_provider(config_type=DataRobotLLMComponentModelConfig)
+async def datarobot_llm_component(
+    config: DataRobotLLMComponentModelConfig, _builder: Builder
+) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot LLM Component for use with an LLM client."
+    )
+
+
+class DataRobotLLMGatewayModelConfig(OpenAIModelConfig, name="datarobot-llm-gateway"):  # type: ignore[call-arg]
+    """A DataRobot LLM provider to be used with an LLM client."""
+
+    api_key: str | None = Field(
+        default=config.datarobot_api_token, description="DataRobot API key."
+    )
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint.rstrip("/"), description="DataRobot LLM gateway URL."
+    )
+
+
+@register_llm_provider(config_type=DataRobotLLMGatewayModelConfig)
+async def datarobot_llm_gateway(
+    config: DataRobotLLMGatewayModelConfig, _builder: Builder
+) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot LLM Gateway for use with an LLM client."
+    )
+
+
+class DataRobotLLMDeploymentModelConfig(OpenAIModelConfig, name="datarobot-llm-deployment"):  # type: ignore[call-arg]
+    """A DataRobot LLM provider to be used with an LLM client."""
+
+    api_key: str | None = Field(
+        default=config.datarobot_api_token, description="DataRobot API key."
+    )
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint + f"/deployments/{config.llm_deployment_id}"
+    )
+    model_name: str = Field(
+        validation_alias=AliasChoices("model_name", "model"),
+        serialization_alias="model",
+        description="The model name to pass through to the deployment.",
+        default="datarobot-deployed-llm",
+    )
+
+
+@register_llm_provider(config_type=DataRobotLLMDeploymentModelConfig)
+async def datarobot_llm_deployment(
+    config: DataRobotLLMDeploymentModelConfig, _builder: Builder
+) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot LLM deployment for use with an LLM client."
+    )
+
+
+class DataRobotNIMModelConfig(DataRobotLLMDeploymentModelConfig, name="datarobot-nim"):  # type: ignore[call-arg]
+    """A DataRobot NIM LLM provider to be used with an LLM client."""
+
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint + f"/deployments/{config.nim_deployment_id}"
+    )
+
+
+@register_llm_provider(config_type=DataRobotNIMModelConfig)
+async def datarobot_nim(config: DataRobotNIMModelConfig, _builder: Builder) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot NIM deployment for use with an LLM client."
+    )