datarobot-genai 0.2.24__py3-none-any.whl → 0.2.29__py3-none-any.whl

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/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:

{datarobot_genai-0.2.24.dist-info → datarobot_genai-0.2.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.24
+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.24.dist-info → datarobot_genai-0.2.29.dist-info}/RECORD

@@ -27,22 +27,21 @@ datarobot_genai/drmcp/server.py,sha256=KE4kjS5f9bfdYftG14HBHrfvxDfCD4pwCXePfvl1O
 datarobot_genai/drmcp/core/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/auth.py,sha256=E-5wrGbBFEBlD5377g6Exddrc7HsazamwX8tWr2RLXY,5815
 datarobot_genai/drmcp/core/clients.py,sha256=y-yG8617LbmiZ_L7FWfMrk4WjIekyr76u_Q80aLqGpI,5524
-datarobot_genai/drmcp/core/config.py,sha256=69QDsVVSvjzv1uIHOjtQGzdg7_Ic4sA3vLA6ZbZJ1Ok,12674
+datarobot_genai/drmcp/core/config.py,sha256=SWLhVKoqI4vmA-04TFKpKm1_G2yEMEN1e_8cv8d_XRM,13774
 datarobot_genai/drmcp/core/config_utils.py,sha256=U-aieWw7MyP03cGDFIp97JH99ZUfr3vD9uuTzBzxn7w,6428
 datarobot_genai/drmcp/core/constants.py,sha256=lUwoW_PTrbaBGqRJifKqCn3EoFacoEgdO-CpoFVrUoU,739
 datarobot_genai/drmcp/core/credentials.py,sha256=PYEUDNMVw1BoMzZKLkPVTypNkVevEPtmk3scKnE-zYg,6706
-datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=yKaIe7Qq23Zgny7Q1dc48iEcpfM8z2Ne6iouGYL58AE,14392
+datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=czcjbwhZAeW9EtG_Bys0GARPOuQulstkiU7FG48Q9bg,14118
 datarobot_genai/drmcp/core/dr_mcp_server_logo.py,sha256=hib-nfR1SNTW6CnpFsFCkL9H_OMwa4YYyinV7VNOuLk,4708
 datarobot_genai/drmcp/core/exceptions.py,sha256=eqsGI-lxybgvWL5w4BFhbm3XzH1eU5tetwjnhJxelpc,905
 datarobot_genai/drmcp/core/logging.py,sha256=Y_hig4eBWiXGaVV7B_3wBcaYVRNH4ydptbEQhrP9-mY,3414
-datarobot_genai/drmcp/core/mcp_instance.py,sha256=hArS-BIdsIdRyRA21a4_ILgqqzmuRxZts-Ewgtf1H60,20917
-datarobot_genai/drmcp/core/mcp_server_tools.py,sha256=odNZKozfx0VV38SLZHw9lY0C0JM_JnRI06W3BBXnyE4,4278
+datarobot_genai/drmcp/core/mcp_instance.py,sha256=nt4gOlAQklMcqmohRIKovYcyhgLdb08NHMo28DBYmOk,18362
 datarobot_genai/drmcp/core/routes.py,sha256=dqE2M0UzAyyN9vQjlyTjYW4rpju3LT039po5weuO__I,17936
 datarobot_genai/drmcp/core/routes_utils.py,sha256=vSseXWlplMSnRgoJgtP_rHxWSAVYcx_tpTv4lyTpQoc,944
 datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_9PgpZT0uyjlv_sr5wc,3419
 datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
-datarobot_genai/drmcp/core/tool_config.py,sha256=5JCWO70ZH-K-34yS7vYJG2nl4i9UO_q_W9NCoWSXXno,3271
-datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
+datarobot_genai/drmcp/core/tool_config.py,sha256=izUdM6dN3GRBzSBs-OagggM2dX5PGBnDbVv4N5bfWFI,3668
+datarobot_genai/drmcp/core/tool_filter.py,sha256=yKQlEtzyIeXGxZJkHbK36QI19vmgQkvqmfx5cTo2pp4,3156
 datarobot_genai/drmcp/core/utils.py,sha256=EvfpqKZ3tECMoxpIQ_tA_3rOgy6KJEYKC0lWZo_Daag,4517
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=AGJlKqgHRO0Kd7Gl-Ulw9KYBgzjTTFXWBvOUF-SuKUI,5454
@@ -73,30 +72,33 @@ datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=46rH0fYYmUj7ygf968iRbdS
 datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=sHA_BWtpgIAFp9IXiNkUeBartBMjLAauqkV9bYtCr-g,3874
 datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=YgyqHK09MB-PBaqT34heqvmvYYFtLpzzSJt7xuTJmDg,11224
 datarobot_genai/drmcp/test_utils/test_interactive.py,sha256=guXvR8q2H6VUdmvIjEJcElQJCC6lQ-oTrzbD2EkHeCs,8025
-datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=wmI-xcL0rSr56-ZoGNB8np0CZHAU583o8-Kw7fRwtMw,6232
+datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=3yMfOsz3LdHYywuE5BhdJDpTUowx37HsFSsMdBTxA80,9337
 datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
 datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
 datarobot_genai/drmcp/tools/clients/confluence.py,sha256=h_G0By_kDnJeWDT_d-IREsaZ5-0xB5GoLXOqblYP5MA,20706
-datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=e28XwX0C8E3nql85-_NbUEMB-4s0lsQ2f5spj9BgsgM,21455
+datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=8GztWTdpJ7Ir3NIvIoOHPzDscoR1Ui7Ct2IiKmuUzIc,26012
 datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
+datarobot_genai/drmcp/tools/clients/microsoft_graph.py,sha256=PASGThDPE8zkBZqach8lurJL1y47DWUPLwvf9N6uLGM,19234
 datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
 datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/confluence/tools.py,sha256=_-ws65WLK8KZP_mKkf4yJ7ZunR8qdyoiMwHQX47MSMw,12362
 datarobot_genai/drmcp/tools/gdrive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=EvoEr3AEI-xRldwCTZHiQfBRHQfLtgHuojTx8mXhlU4,7074
+datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=G8LlnGEINZqV83Q-b3ZliWkDjouhbozDam3w6GfA7s0,10711
 datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
+datarobot_genai/drmcp/tools/microsoft_graph/__init__.py,sha256=CuOaMt1AJo7cHx_GuhO3s_aqxZas_wlDsoBorBsvbeU,577
+datarobot_genai/drmcp/tools/microsoft_graph/tools.py,sha256=zJ-UA1TMhPOYcExvgWv0YBjDsSIDPA-U1SEbBrVfAc8,7744
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
-datarobot_genai/drmcp/tools/predictive/data.py,sha256=mt3PvIel4IUAZo0HLIs6QsTtRkuzX1qg-5PHI3IJ7E8,4455
+datarobot_genai/drmcp/tools/predictive/data.py,sha256=sSFAmO6x0DSuolw8urhMaOj5PwfUH29oc2mEOZI3YU4,4631
 datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
 datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
 datarobot_genai/drmcp/tools/predictive/model.py,sha256=Yih5-KedJ-1yupPLXCJsCXOdyWWi9pRvgapXDlgXWJA,4891
 datarobot_genai/drmcp/tools/predictive/predict.py,sha256=Qoob2_t2crfWtyPzkXMRz2ITZumnczU6Dq4C7q9RBMI,9370
 datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=urq6rPyZFsAP-bPyclSNzrkvb6FTamdlFau8q0IWWJ0,13472
-datarobot_genai/drmcp/tools/predictive/project.py,sha256=KaMDAvJY4s12j_4ybA7-KcCS1yMOj-KPIKNBgCSE2iM,2536
-datarobot_genai/drmcp/tools/predictive/training.py,sha256=S9V7AlO6mAgIAJNww0g5agFOw4YqRiCsIGaRDJcOe4A,23991
+datarobot_genai/drmcp/tools/predictive/project.py,sha256=xC52UdYvuFeNZC7Y5MfXcvzTL70WwAacQXESr6rqN6s,3255
+datarobot_genai/drmcp/tools/predictive/training.py,sha256=LzMxbBT8wxKYDrRlVElfmTUrzpmGvwrR-mTGf6YUnIA,23998
 datarobot_genai/langgraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/langgraph/agent.py,sha256=DRnywmS9KDywyChtuIZZwNKbJs8BpC259EG_kxYbiQ8,15828
 datarobot_genai/langgraph/mcp.py,sha256=iA2_j46mZAaNaL7ntXT-LW6C-NMJkzr3VfKDDfe7mh8,2851
@@ -107,13 +109,13 @@ datarobot_genai/llama_index/mcp.py,sha256=leXqF1C4zhuYEKFwNEfZHY4dsUuGZk3W7KArY-
 datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/nat/agent.py,sha256=DuGrgqt1FzvAE-cRH_P3LTFUlwuClvbVurdwA-RsbuY,11177
 datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZC74ziNo_FHbsZgtiM,4230
-datarobot_genai/nat/datarobot_llm_clients.py,sha256=Yu208Ed_p_4P3HdpuM7fYnKcXtimORHpKlWVPyijpU8,11356
+datarobot_genai/nat/datarobot_llm_clients.py,sha256=-_q_KlKOVQecIYJd8YRiYnS4ZNazQAiAdZBE1Zip_wQ,12684
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
-datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
+datarobot_genai/nat/datarobot_mcp_client.py,sha256=jL8sXb8g4gvt0VYgB2tfMGsMjpB1GV2XIbN0iv_LxVU,10701
 datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
-datarobot_genai-0.2.24.dist-info/METADATA,sha256=-QGxEqh8oSYn-gL93kIxqGhOTkrlHQRMf4HqMYTyfDs,6301
-datarobot_genai-0.2.24.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.24.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.24.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.24.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.24.dist-info/RECORD,,
+datarobot_genai-0.2.29.dist-info/METADATA,sha256=lFk85PaEbHw_waws2INyPqFxo92s43jFp4vFyz8HHdc,6301
+datarobot_genai-0.2.29.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.29.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.29.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.29.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.29.dist-info/RECORD,,

datarobot_genai/drmcp/core/mcp_server_tools.py

@@ -1,129 +0,0 @@
-# Copyright 2025 DataRobot, Inc.
-#
-# 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.
-
-import logging
-
-from .mcp_instance import dr_core_mcp_tool
-from .mcp_instance import mcp
-
-logger = logging.getLogger(__name__)
-
-
-@dr_core_mcp_tool(tags={"mcp_server_tools", "metadata"})
-async def get_all_available_tags() -> str:
-    """
-    List all unique tags from all registered tools.
-
-    Returns
-    -------
-        A string containing all available tags, one per line.
-    """
-    tags = await mcp.get_all_tags()
-    if not tags:
-        return "No tags found in any tools."
-
-    return "\n".join(sorted(tags))
-
-
-@dr_core_mcp_tool(tags={"mcp_server_tools", "metadata", "discovery"})
-async def list_tools_by_tags(tags: list[str] | None = None, match_all: bool = False) -> str:
-    """
-    List tools filtered by tags.
-
-    Args:
-        tags: Optional list of tags to filter by. If None, returns all tools.
-        match_all: If True, tool must have all specified tags (AND logic).
-                  If False, tool must have at least one tag (OR logic).
-                  Only used when tags is provided.
-
-    Returns
-    -------
-        A formatted string listing tools that match the tag criteria.
-    """
-    tools = await mcp.list_tools(tags=tags, match_all=match_all)
-
-    if not tools:
-        if tags:
-            logic = "all" if match_all else "any"
-            return f"No tools found with {logic} of the tags: {', '.join(tags)}"
-        else:
-            return "No tools found."
-
-    result = []
-    if tags:
-        logic = "all" if match_all else "any"
-        result.append(f"Tools with {logic} of the tags: {', '.join(tags)}")
-    else:
-        result.append("All available tools:")
-
-    result.append("")
-
-    for i, tool in enumerate(tools, 1):
-        tool_tags = []
-        if tool.annotations and hasattr(tool.annotations, "extra") and tool.annotations.extra:
-            tool_tags = tool.annotations.extra.get("tags", [])
-
-        result.append(f"{i}. {tool.name}")
-        result.append(f"   Description: {tool.description}")
-        if tool_tags:
-            result.append(f"   Tags: {', '.join(tool_tags)}")
-        result.append("")
-
-    return "\n".join(result)
-
-
-@dr_core_mcp_tool(tags={"mcp_server_tools", "metadata", "discovery"})
-async def get_tool_info_by_name(tool_name: str) -> str:
-    """
-    Get detailed information about a specific tool by name.
-
-    Args:
-        tool_name: The name of the tool to get information about.
-
-    Returns
-    -------
-        A formatted string with detailed information about the tool.
-    """
-    all_tools = await mcp.list_tools()
-
-    for tool in all_tools:
-        if tool.name == tool_name:
-            result = [f"Tool: {tool.name}"]
-            result.append(f"Description: {tool.description}")
-
-            # Get tags
-            tool_tags = []
-            if tool.annotations and hasattr(tool.annotations, "extra") and tool.annotations.extra:
-                tool_tags = tool.annotations.extra.get("tags", [])
-
-            if tool_tags:
-                result.append(f"Tags: {', '.join(tool_tags)}")
-            else:
-                result.append("Tags: None")
-
-            # Get input schema info
-            if (
-                tool.inputSchema
-                and hasattr(tool.inputSchema, "properties")
-                and tool.inputSchema.properties
-            ):
-                result.append("Parameters:")
-                for param_name, param_info in tool.inputSchema.properties.items():
-                    param_type = param_info.get("type", "unknown")
-                    param_desc = param_info.get("description", "No description")
-                    result.append(f"  - {param_name} ({param_type}): {param_desc}")
-
-            return "\n".join(result)
-
-    return f"Tool '{tool_name}' not found."