datarobot-genai 0.2.28__tar.gz → 0.2.29__tar.gz

{datarobot_genai-0.2.28 → datarobot_genai-0.2.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.28
+Version: 0.2.29
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.2.28 → datarobot_genai-0.2.29}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.28"
+version = "0.2.29"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.2.28 → datarobot_genai-0.2.29}/src/datarobot_genai/nat/datarobot_llm_clients.py

@@ -12,22 +12,18 @@
 # 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 crewai import LLM
-from langchain_openai import ChatOpenAI
-from llama_index.core.base.llms.types import LLMMetadata
-from llama_index.llms.litellm import LiteLLM
 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.plugins.langchain.llm import (
-    _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
-)
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 
 from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
@@ -35,6 +31,11 @@ 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")
 
 
@@ -50,42 +51,53 @@ def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) ->
     return client
 
 
-class DataRobotChatOpenAI(ChatOpenAI):
-    def _get_request_payload(
-        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
-
-
-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.
+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.
         """
-        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,
-        )
+
+        @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(
@@ -94,11 +106,15 @@ class DataRobotLiteLLM(LiteLLM):  # type: ignore[misc]
 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 = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -108,6 +124,8 @@ async def datarobot_llm_gateway_langchain(
 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"]
@@ -121,12 +139,12 @@ async def datarobot_llm_gateway_crewai(
 )
 async def datarobot_llm_gateway_llamaindex(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> 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 = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -136,6 +154,10 @@ async def datarobot_llm_gateway_llamaindex(
 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,
@@ -143,7 +165,7 @@ async def datarobot_llm_deployment_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -153,6 +175,8 @@ async def datarobot_llm_deployment_langchain(
 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,
@@ -170,7 +194,7 @@ async def datarobot_llm_deployment_crewai(
 )
 async def datarobot_llm_deployment_llamaindex(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -179,7 +203,7 @@ async def datarobot_llm_deployment_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -187,6 +211,10 @@ async def datarobot_llm_deployment_llamaindex(
 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,
@@ -194,7 +222,7 @@ async def datarobot_nim_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    client = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -202,6 +230,8 @@ async def datarobot_nim_langchain(
 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,
@@ -217,7 +247,7 @@ async def datarobot_nim_crewai(
 @register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
 async def datarobot_nim_llamaindex(
     llm_config: DataRobotNIMModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> AsyncGenerator[LiteLLM]:
     config = llm_config.model_dump(
         exclude={"type", "thinking"},
         by_alias=True,
@@ -226,7 +256,7 @@ async def datarobot_nim_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)
 
 
@@ -236,13 +266,17 @@ async def datarobot_nim_llamaindex(
 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 = DataRobotChatOpenAI(**config)
+    client = _create_datarobot_chat_openai(config)
     yield langchain_patch_llm_based_on_config(client, config)
 
 
@@ -252,6 +286,8 @@ async def datarobot_llm_component_langchain(
 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"]
@@ -269,7 +305,7 @@ async def datarobot_llm_component_crewai(
 )
 async def datarobot_llm_component_llamaindex(
     llm_config: DataRobotLLMComponentModelConfig, builder: Builder
-) -> AsyncGenerator[LLM]:
+) -> 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"]
@@ -278,5 +314,5 @@ async def datarobot_llm_component_llamaindex(
     else:
         config["api_base"] = config.pop("base_url") + "/chat/completions"
     config.pop("use_datarobot_llm_gateway")
-    client = DataRobotLiteLLM(**config)
+    client = _create_datarobot_litellm(config)
     yield _patch_llm_based_on_config(client, config)

{datarobot_genai-0.2.28 → datarobot_genai-0.2.29}/src/datarobot_genai/nat/datarobot_mcp_client.py

@@ -12,57 +12,83 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 from datetime import timedelta
+from typing import TYPE_CHECKING
 from typing import Literal
 
-import httpx
-from nat.authentication.interfaces import AuthProviderBase
-from nat.builder.builder import Builder
 from nat.cli.register_workflow import register_function_group
 from nat.data_models.component_ref import AuthenticationRef
 from nat.plugins.mcp.client_base import AuthAdapter
-from nat.plugins.mcp.client_base import MCPSSEClient
-from nat.plugins.mcp.client_base import MCPStdioClient
 from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
 from nat.plugins.mcp.client_config import MCPServerConfig
 from nat.plugins.mcp.client_impl import MCPClientConfig
-from nat.plugins.mcp.client_impl import MCPFunctionGroup
-from nat.plugins.mcp.client_impl import mcp_apply_tool_alias_and_description
-from nat.plugins.mcp.client_impl import mcp_session_tool_function
 from pydantic import Field
 from pydantic import HttpUrl
 
-from datarobot_genai.core.mcp.common import MCPConfig
+if TYPE_CHECKING:
+    import httpx
+    from nat.authentication.interfaces import AuthProviderBase
+    from nat.builder.builder import Builder
+    from nat.plugins.mcp.client_impl import MCPFunctionGroup
 
 logger = logging.getLogger(__name__)
 
-config = MCPConfig().server_config
+
+def _default_transport() -> Literal["streamable-http", "sse", "stdio"]:
+    from datarobot_genai.core.mcp.common import MCPConfig  # noqa: PLC0415
+
+    server_config = MCPConfig().server_config
+    return server_config["transport"] if server_config else "stdio"
+
+
+def _default_url() -> HttpUrl | None:
+    from datarobot_genai.core.mcp.common import MCPConfig  # noqa: PLC0415
+
+    server_config = MCPConfig().server_config
+    return server_config["url"] if server_config else None
+
+
+def _default_auth_provider() -> str | AuthenticationRef | None:
+    from datarobot_genai.core.mcp.common import MCPConfig  # noqa: PLC0415
+
+    server_config = MCPConfig().server_config
+    return "datarobot_mcp_auth" if server_config else None
+
+
+def _default_command() -> str | None:
+    from datarobot_genai.core.mcp.common import MCPConfig  # noqa: PLC0415
+
+    server_config = MCPConfig().server_config
+    return None if server_config else "docker"
 
 
 class DataRobotMCPServerConfig(MCPServerConfig):
     transport: Literal["streamable-http", "sse", "stdio"] = Field(
-        default=config["transport"] if config else "stdio",
+        default_factory=_default_transport,
         description="Transport type to connect to the MCP server (sse or streamable-http)",
     )
     url: HttpUrl | None = Field(
-        default=config["url"] if config else None,
+        default_factory=_default_url,
         description="URL of the MCP server (for sse or streamable-http transport)",
     )
     # Authentication configuration
     auth_provider: str | AuthenticationRef | None = Field(
-        default="datarobot_mcp_auth" if config else None,
+        default_factory=_default_auth_provider,
         description="Reference to authentication provider",
     )
     command: str | None = Field(
-        default=None if config else "docker",
+        default_factory=_default_command,
         description="Command to run for stdio transport (e.g. 'python' or 'docker')",
     )
 
 
 class DataRobotMCPClientConfig(MCPClientConfig, name="datarobot_mcp_client"):  # type: ignore[call-arg]
     server: DataRobotMCPServerConfig = Field(
-        default=DataRobotMCPServerConfig(), description="DataRobot MCP Server configuration"
+        default_factory=DataRobotMCPServerConfig,
+        description="DataRobot MCP Server configuration",
     )
 
 
@@ -128,6 +154,12 @@ async def datarobot_mcp_client_function_group(
     Returns:
         The function group
     """
+    from nat.plugins.mcp.client_base import MCPSSEClient  # noqa: PLC0415
+    from nat.plugins.mcp.client_base import MCPStdioClient  # noqa: PLC0415
+    from nat.plugins.mcp.client_impl import MCPFunctionGroup  # noqa: PLC0415
+    from nat.plugins.mcp.client_impl import mcp_apply_tool_alias_and_description  # noqa: PLC0415
+    from nat.plugins.mcp.client_impl import mcp_session_tool_function  # noqa: PLC0415
+
     # Resolve auth provider if specified
     auth_provider = None
     if config.server.auth_provider: