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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.28
+Version: 0.2.30
 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.30}/pyproject.toml

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

{datarobot_genai-0.2.28 → datarobot_genai-0.2.30}/src/datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -33,7 +33,17 @@ from datarobot_genai.drmcp.core.auth import get_access_token
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_FIELDS = {"id", "name", "size", "mimeType", "webViewLink", "createdTime", "modifiedTime"}
+SUPPORTED_FIELDS = {
+    "id",
+    "name",
+    "size",
+    "mimeType",
+    "webViewLink",
+    "createdTime",
+    "modifiedTime",
+    "starred",
+    "trashed",
+}
 SUPPORTED_FIELDS_STR = ",".join(SUPPORTED_FIELDS)
 DEFAULT_FIELDS = f"nextPageToken,files({SUPPORTED_FIELDS_STR})"
 GOOGLE_DRIVE_FOLDER_MIME = "application/vnd.google-apps.folder"
@@ -119,6 +129,8 @@ class GoogleDriveFile(BaseModel):
     web_view_link: Annotated[str | None, Field(alias="webViewLink")] = None
     created_time: Annotated[str | None, Field(alias="createdTime")] = None
     modified_time: Annotated[str | None, Field(alias="modifiedTime")] = None
+    starred: bool | None = None
+    trashed: bool | None = None
 
     model_config = ConfigDict(populate_by_name=True)
 
@@ -133,8 +145,31 @@ class GoogleDriveFile(BaseModel):
             web_view_link=data.get("webViewLink"),
             created_time=data.get("createdTime"),
             modified_time=data.get("modifiedTime"),
+            starred=data.get("starred"),
+            trashed=data.get("trashed"),
         )
 
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the file."""
+        result: dict[str, Any] = {
+            "id": self.id,
+            "name": self.name,
+            "mimeType": self.mime_type,
+        }
+        if self.size is not None:
+            result["size"] = self.size
+        if self.web_view_link is not None:
+            result["webViewLink"] = self.web_view_link
+        if self.created_time is not None:
+            result["createdTime"] = self.created_time
+        if self.modified_time is not None:
+            result["modifiedTime"] = self.modified_time
+        if self.starred is not None:
+            result["starred"] = self.starred
+        if self.trashed is not None:
+            result["trashed"] = self.trashed
+        return result
+
 
 class PaginatedResult(BaseModel):
     """Result of a paginated API call."""
@@ -440,6 +475,66 @@ class GoogleDriveClient:
         response.raise_for_status()
         return GoogleDriveFile.from_api_response(response.json())
 
+    async def update_file_metadata(
+        self,
+        file_id: str,
+        new_name: str | None = None,
+        starred: bool | None = None,
+        trashed: bool | None = None,
+    ) -> GoogleDriveFile:
+        """Update file metadata in Google Drive.
+
+        Args:
+            file_id: The ID of the file to update.
+            new_name: A new name to rename the file. Must not be empty or whitespace.
+            starred: Set to True to star the file or False to unstar it.
+            trashed: Set to True to trash the file or False to restore it.
+
+        Returns
+        -------
+            GoogleDriveFile with updated metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If no update fields are provided, file is not found,
+                             access is denied, or the request is invalid.
+        """
+        if new_name is None and starred is None and trashed is None:
+            raise GoogleDriveError(
+                "At least one of new_name, starred, or trashed must be provided."
+            )
+
+        if new_name is not None and not new_name.strip():
+            raise GoogleDriveError("new_name cannot be empty or whitespace.")
+
+        body: dict[str, Any] = {}
+        if new_name is not None:
+            body["name"] = new_name
+        if starred is not None:
+            body["starred"] = starred
+        if trashed is not None:
+            body["trashed"] = trashed
+
+        response = await self._client.patch(
+            f"/{file_id}",
+            json=body,
+            params={"fields": SUPPORTED_FIELDS_STR, "supportsAllDrives": "true"},
+        )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(f"File with ID '{file_id}' not found.")
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                f"Permission denied: you don't have permission to update file '{file_id}'."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError("Bad request: invalid parameters for file update.")
+        if response.status_code == 429:
+            raise GoogleDriveError("Rate limit exceeded. Please try again later.")
+
+        response.raise_for_status()
+        return GoogleDriveFile.from_api_response(response.json())
+
     async def _export_workspace_file(self, file_id: str, export_mime_type: str) -> str:
         """Export a Google Workspace file to the specified format.
 

{datarobot_genai-0.2.28 → datarobot_genai-0.2.30}/src/datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -164,7 +164,6 @@ async def gdrive_read_content(
             f"An unexpected error occurred while reading Google Drive file content: {str(e)}"
         )
 
-    # Provide helpful context about the conversion
     export_info = ""
     if file_content.was_exported:
         export_info = f" (exported from {file_content.original_mime_type})"
@@ -252,7 +251,6 @@ async def gdrive_create_file(
         logger.error(f"Unexpected error creating Google Drive file: {e}")
         raise ToolError(f"An unexpected error occurred while creating Google Drive file: {str(e)}")
 
-    # Build response message
     file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
     content_info = ""
     if initial_content and mime_type != GOOGLE_DRIVE_FOLDER_MIME:
@@ -260,11 +258,90 @@ async def gdrive_create_file(
 
     return ToolResult(
         content=f"Successfully created {file_type} '{created_file.name}'{content_info}.",
-        structured_content={
-            "id": created_file.id,
-            "name": created_file.name,
-            "mimeType": created_file.mime_type,
-            "webViewLink": created_file.web_view_link,
-            "createdTime": created_file.created_time,
-        },
+        structured_content=created_file.as_flat_dict(),
+    )
+
+
+@dr_mcp_tool(
+    tags={"google", "gdrive", "update", "metadata", "rename", "star", "trash"}, enabled=False
+)
+async def gdrive_update_metadata(
+    *,
+    file_id: Annotated[str, "The ID of the file or folder to update."],
+    new_name: Annotated[str | None, "A new name to rename the file."] = None,
+    starred: Annotated[bool | None, "Set to True to star the file or False to unstar it."] = None,
+    trash: Annotated[bool | None, "Set to True to trash the file or False to restore it."] = None,
+) -> ToolResult:
+    """
+    Update non-content metadata fields of a Google Drive file or folder.
+
+    This tool allows you to:
+    - Rename files and folders by setting new_name
+    - Star or unstar files (per-user preference) with starred
+    - Move files to trash or restore them with trash
+
+    Usage:
+        - Rename: gdrive_update_metadata(file_id="1ABC...", new_name="New Name.txt")
+        - Star: gdrive_update_metadata(file_id="1ABC...", starred=True)
+        - Unstar: gdrive_update_metadata(file_id="1ABC...", starred=False)
+        - Trash: gdrive_update_metadata(file_id="1ABC...", trash=True)
+        - Restore: gdrive_update_metadata(file_id="1ABC...", trash=False)
+        - Multiple: gdrive_update_metadata(file_id="1ABC...", new_name="New.txt", starred=True)
+
+    Note:
+        - At least one of new_name, starred, or trash must be provided.
+        - Starring is per-user: starring a shared file only affects your view.
+        - Trashing a folder trashes all contents recursively.
+        - Trashing requires permissions (owner for My Drive, organizer for Shared Drives).
+    """
+    if not file_id or not file_id.strip():
+        raise ToolError("Argument validation error: 'file_id' cannot be empty.")
+
+    if new_name is None and starred is None and trash is None:
+        raise ToolError(
+            "Argument validation error: at least one of 'new_name', 'starred', or 'trash' "
+            "must be provided."
+        )
+
+    if new_name is not None and not new_name.strip():
+        raise ToolError("Argument validation error: 'new_name' cannot be empty or whitespace.")
+
+    access_token = await get_gdrive_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with GoogleDriveClient(access_token) as client:
+            updated_file = await client.update_file_metadata(
+                file_id=file_id,
+                new_name=new_name,
+                starred=starred,
+                trashed=trash,
+            )
+    except GoogleDriveError as e:
+        logger.error(f"Google Drive error updating file metadata: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error updating Google Drive file metadata: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while updating Google Drive file metadata: {str(e)}"
+        )
+
+    changes: list[str] = []
+    if new_name is not None:
+        changes.append(f"renamed to '{new_name}'")
+    if starred is True:
+        changes.append("starred")
+    elif starred is False:
+        changes.append("unstarred")
+    if trash is True:
+        changes.append("moved to trash")
+    elif trash is False:
+        changes.append("restored from trash")
+
+    changes_description = ", ".join(changes)
+
+    return ToolResult(
+        content=f"Successfully updated file '{updated_file.name}': {changes_description}.",
+        structured_content=updated_file.as_flat_dict(),
     )

{datarobot_genai-0.2.28 → datarobot_genai-0.2.30}/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.30}/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: