datarobot-genai 0.1.75__py3-none-any.whl → 0.2.11__py3-none-any.whl

datarobot_genai/drmcp/tools/jira/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+

datarobot_genai/drmcp/tools/jira/tools.py

@@ -0,0 +1,52 @@
+# 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 typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.atlassian import get_atlassian_access_token
+from datarobot_genai.drmcp.tools.clients.jira import JiraClient
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"jira", "read", "get", "issue"})
+async def jira_get_issue(
+    *, issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."]
+) -> ToolResult:
+    """Retrieve all fields and details for a single Jira issue by its key."""
+    if not issue_key:
+        raise ToolError("Argument validation error: 'issue_key' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with JiraClient(access_token) as client:
+            issue = await client.get_jira_issue(issue_key)
+    except Exception as e:
+        logger.error(f"Unexpected error getting Jira issue: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Successfully retrieved details for issue '{issue_key}'.",
+        structured_content=issue.as_flat_dict(),
+    )

datarobot_genai/drmcp/tools/predictive/predict.py

@@ -22,10 +22,10 @@ from fastmcp.resources import HttpResource
 from fastmcp.resources import ResourceManager
 
 from datarobot_genai.drmcp.core.clients import get_credentials
-from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 from datarobot_genai.drmcp.core.utils import generate_presigned_url
+from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
 
 logger = logging.getLogger(__name__)
 

datarobot_genai/drmcp/tools/predictive/predict_realtime.py

@@ -23,11 +23,11 @@ from datarobot_predict import TimeSeriesType
 from datarobot_predict.deployment import predict as dr_predict
 from pydantic import BaseModel
 
-from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 from datarobot_genai.drmcp.core.utils import PredictionResponse
 from datarobot_genai.drmcp.core.utils import predictions_result_response
+from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
 
 logger = logging.getLogger(__name__)
 

datarobot_genai/langgraph/agent.py

@@ -17,6 +17,15 @@ from collections.abc import AsyncGenerator
 from typing import Any
 from typing import cast
 
+from ag_ui.core import Event
+from ag_ui.core import EventType
+from ag_ui.core import TextMessageContentEvent
+from ag_ui.core import TextMessageEndEvent
+from ag_ui.core import TextMessageStartEvent
+from ag_ui.core import ToolCallArgsEvent
+from ag_ui.core import ToolCallEndEvent
+from ag_ui.core import ToolCallResultEvent
+from ag_ui.core import ToolCallStartEvent
 from langchain.tools import BaseTool
 from langchain_core.messages import AIMessageChunk
 from langchain_core.messages import ToolMessage
@@ -158,43 +167,7 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
         # The main difference is returning a generator for streaming or a final response for sync.
         if is_streaming(completion_create_params):
             # Streaming response: yield each message as it is generated
-            async def stream_generator() -> AsyncGenerator[
-                tuple[str, MultiTurnSample | None, UsageMetrics], None
-            ]:
-                # Iterate over the graph stream. For message events, yield the content.
-                # For update events, accumulate the usage metrics.
-                events = []
-                async for _, mode, event in graph_stream:
-                    if mode == "messages":
-                        message_event: tuple[AIMessageChunk, dict[str, Any]] = event  # type: ignore[assignment]
-                        llm_token, _ = message_event
-                        yield (
-                            str(llm_token.content),
-                            None,
-                            usage_metrics,
-                        )
-                    elif mode == "updates":
-                        update_event: dict[str, Any] = event  # type: ignore[assignment]
-                        events.append(update_event)
-                        current_node = next(iter(update_event))
-                        node_data = update_event[current_node]
-                        current_usage = node_data.get("usage", {}) if node_data is not None else {}
-                        if current_usage:
-                            usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
-                            usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
-                            usage_metrics["completion_tokens"] += current_usage.get(
-                                "completion_tokens", 0
-                            )
-                    else:
-                        raise ValueError(f"Invalid mode: {mode}")
-
-                # Create a list of events from the event listener
-                pipeline_interactions = self.create_pipeline_interactions_from_events(events)
-
-                # yield the final response indicating completion
-                yield "", pipeline_interactions, usage_metrics
-
-            return stream_generator()
+            return self._stream_generator(graph_stream, usage_metrics)
         else:
             # Synchronous response: collect all events and return the final message
             events: list[dict[str, Any]] = [
@@ -203,6 +176,16 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
                 if mode == "updates"
             ]
 
+            # Accumulate the usage metrics from the updates
+            for update in events:
+                current_node = next(iter(update))
+                node_data = update[current_node]
+                current_usage = node_data.get("usage", {}) if node_data is not None else {}
+                if current_usage:
+                    usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
+                    usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
+                    usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
+
             pipeline_interactions = self.create_pipeline_interactions_from_events(events)
 
             # Extract the final event from the graph stream as the synchronous response
@@ -214,14 +197,132 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
                 if node_data is not None and "messages" in node_data
                 else ""
             )
-            current_usage = node_data.get("usage", {}) if node_data is not None else {}
-            if current_usage:
-                usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
-                usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
-                usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
 
             return response_text, pipeline_interactions, usage_metrics
 
+    async def _stream_generator(
+        self, graph_stream: AsyncGenerator[tuple[Any, str, Any], None], usage_metrics: UsageMetrics
+    ) -> AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]:
+        # Iterate over the graph stream. For message events, yield the content.
+        # For update events, accumulate the usage metrics.
+        events = []
+        current_message_id = None
+        tool_call_id = ""
+        async for _, mode, event in graph_stream:
+            if mode == "messages":
+                message_event: tuple[AIMessageChunk | ToolMessage, dict[str, Any]] = event  # type: ignore[assignment]
+                message = message_event[0]
+                if isinstance(message, ToolMessage):
+                    yield (
+                        ToolCallEndEvent(
+                            type=EventType.TOOL_CALL_END, tool_call_id=message.tool_call_id
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    yield (
+                        ToolCallResultEvent(
+                            type=EventType.TOOL_CALL_RESULT,
+                            message_id=message.id,
+                            tool_call_id=message.tool_call_id,
+                            content=message.content,
+                            role="tool",
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    tool_call_id = ""
+                elif isinstance(message, AIMessageChunk):
+                    if message.tool_call_chunks:
+                        # This is a tool call message
+                        for tool_call_chunk in message.tool_call_chunks:
+                            if name := tool_call_chunk.get("name"):
+                                # Its a tool call start message
+                                tool_call_id = tool_call_chunk["id"]
+                                yield (
+                                    ToolCallStartEvent(
+                                        type=EventType.TOOL_CALL_START,
+                                        tool_call_id=tool_call_id,
+                                        tool_call_name=name,
+                                        parent_message_id=message.id,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                            elif args := tool_call_chunk.get("args"):
+                                # Its a tool call args message
+                                yield (
+                                    ToolCallArgsEvent(
+                                        type=EventType.TOOL_CALL_ARGS,
+                                        # Its empty when the tool chunk is not a start message
+                                        # So we use the tool call id from a previous start message
+                                        tool_call_id=tool_call_id,
+                                        delta=args,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                    elif message.content:
+                        # Its a text message
+                        # Handle the start and end of the text message
+                        if message.id != current_message_id:
+                            if current_message_id:
+                                yield (
+                                    TextMessageEndEvent(
+                                        type=EventType.TEXT_MESSAGE_END,
+                                        message_id=current_message_id,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                            current_message_id = message.id
+                            yield (
+                                TextMessageStartEvent(
+                                    type=EventType.TEXT_MESSAGE_START,
+                                    message_id=message.id,
+                                    role="assistant",
+                                ),
+                                None,
+                                usage_metrics,
+                            )
+                        yield (
+                            TextMessageContentEvent(
+                                type=EventType.TEXT_MESSAGE_CONTENT,
+                                message_id=message.id,
+                                delta=message.content,
+                            ),
+                            None,
+                            usage_metrics,
+                        )
+                else:
+                    raise ValueError(f"Invalid message event: {message_event}")
+            elif mode == "updates":
+                update_event: dict[str, Any] = event  # type: ignore[assignment]
+                events.append(update_event)
+                current_node = next(iter(update_event))
+                node_data = update_event[current_node]
+                current_usage = node_data.get("usage", {}) if node_data is not None else {}
+                if current_usage:
+                    usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
+                    usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
+                    usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
+                if current_message_id:
+                    yield (
+                        TextMessageEndEvent(
+                            type=EventType.TEXT_MESSAGE_END,
+                            message_id=current_message_id,
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    current_message_id = None
+
+        # Create a list of events from the event listener
+        pipeline_interactions = self.create_pipeline_interactions_from_events(events)
+
+        # yield the final response indicating completion
+        yield "", pipeline_interactions, usage_metrics
+
     @classmethod
     def create_pipeline_interactions_from_events(
         cls,

datarobot_genai/nat/agent.py

@@ -123,6 +123,8 @@ class NatAgent(BaseAgent[None]):
         model: str | None = None,
         verbose: bool | str | None = True,
         timeout: int | None = 90,
+        authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -131,6 +133,8 @@ class NatAgent(BaseAgent[None]):
             model=model,
             verbose=verbose,
             timeout=timeout,
+            authorization_context=authorization_context,
+            forwarded_headers=forwarded_headers,
             **kwargs,
         )
         self.workflow_path = workflow_path

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_mcp_client.py

@@ -0,0 +1,234 @@
+# 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.
+
+import logging
+from datetime import timedelta
+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
+
+logger = logging.getLogger(__name__)
+
+config = MCPConfig().server_config
+
+
+class DataRobotMCPServerConfig(MCPServerConfig):
+    transport: Literal["streamable-http", "sse", "stdio"] = Field(
+        default=config["transport"] if config else "stdio",
+        description="Transport type to connect to the MCP server (sse or streamable-http)",
+    )
+    url: HttpUrl | None = Field(
+        default=config["url"] if config else None,
+        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,
+        description="Reference to authentication provider",
+    )
+    command: str | None = Field(
+        default=None if config else "docker",
+        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"
+    )
+
+
+class DataRobotAuthAdapter(AuthAdapter):
+    async def _get_auth_headers(
+        self, request: httpx.Request | None = None, response: httpx.Response | None = None
+    ) -> dict[str, str]:
+        """Get authentication headers from the NAT auth provider."""
+        try:
+            # Use the user_id passed to this AuthAdapter instance
+            auth_result = await self.auth_provider.authenticate(
+                user_id=self.user_id, response=response
+            )
+            as_kwargs = auth_result.as_requests_kwargs()
+            return as_kwargs["headers"]
+        except Exception as e:
+            logger.warning("Failed to get auth token: %s", e)
+            return {}
+
+
+class DataRobotMCPStreamableHTTPClient(MCPStreamableHTTPClient):
+    def __init__(
+        self,
+        url: str,
+        auth_provider: AuthProviderBase | None = None,
+        user_id: str | None = None,
+        tool_call_timeout: timedelta = timedelta(seconds=60),
+        auth_flow_timeout: timedelta = timedelta(seconds=300),
+        reconnect_enabled: bool = True,
+        reconnect_max_attempts: int = 2,
+        reconnect_initial_backoff: float = 0.5,
+        reconnect_max_backoff: float = 50.0,
+    ):
+        super().__init__(
+            url=url,
+            auth_provider=auth_provider,
+            user_id=user_id,
+            tool_call_timeout=tool_call_timeout,
+            auth_flow_timeout=auth_flow_timeout,
+            reconnect_enabled=reconnect_enabled,
+            reconnect_max_attempts=reconnect_max_attempts,
+            reconnect_initial_backoff=reconnect_initial_backoff,
+            reconnect_max_backoff=reconnect_max_backoff,
+        )
+        effective_user_id = user_id or (
+            auth_provider.config.default_user_id if auth_provider else None
+        )
+        self._httpx_auth = (
+            DataRobotAuthAdapter(auth_provider, effective_user_id) if auth_provider else None
+        )
+
+
+@register_function_group(config_type=DataRobotMCPClientConfig)
+async def datarobot_mcp_client_function_group(
+    config: DataRobotMCPClientConfig, _builder: Builder
+) -> MCPFunctionGroup:
+    """
+    Connect to an MCP server and expose tools as a function group.
+
+    Args:
+        config: The configuration for the MCP client
+        _builder: The builder
+    Returns:
+        The function group
+    """
+    # Resolve auth provider if specified
+    auth_provider = None
+    if config.server.auth_provider:
+        auth_provider = await _builder.get_auth_provider(config.server.auth_provider)
+
+    # Build the appropriate client
+    if config.server.transport == "stdio":
+        if not config.server.command:
+            raise ValueError("command is required for stdio transport")
+        client = MCPStdioClient(
+            config.server.command,
+            config.server.args,
+            config.server.env,
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    elif config.server.transport == "sse":
+        client = MCPSSEClient(
+            str(config.server.url),
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    elif config.server.transport == "streamable-http":
+        # Use default_user_id for the base client
+        base_user_id = auth_provider.config.default_user_id if auth_provider else None
+        client = DataRobotMCPStreamableHTTPClient(
+            str(config.server.url),
+            auth_provider=auth_provider,
+            user_id=base_user_id,
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    else:
+        raise ValueError(f"Unsupported transport: {config.server.transport}")
+
+    logger.info("Configured to use MCP server at %s", client.server_name)
+
+    # Create the MCP function group
+    group = MCPFunctionGroup(config=config)
+
+    # Store shared components for session client creation
+    group._shared_auth_provider = auth_provider
+    group._client_config = config
+
+    async with client:
+        # Expose the live MCP client on the function group instance so other components
+        # (e.g., HTTP endpoints) can reuse the already-established session instead of creating a
+        # new client per request.
+        group.mcp_client = client
+        group.mcp_client_server_name = client.server_name
+        group.mcp_client_transport = client.transport
+
+        all_tools = await client.get_tools()
+        tool_overrides = mcp_apply_tool_alias_and_description(all_tools, config.tool_overrides)
+
+        # Add each tool as a function to the group
+        for tool_name, tool in all_tools.items():
+            # Get override if it exists
+            override = tool_overrides.get(tool_name)
+
+            # Use override values or defaults
+            function_name = override.alias if override and override.alias else tool_name
+            description = (
+                override.description if override and override.description else tool.description
+            )
+
+            # Create the tool function according to configuration
+            tool_fn = mcp_session_tool_function(tool, group)
+
+            # Normalize optional typing for linter/type-checker compatibility
+            single_fn = tool_fn.single_fn
+            if single_fn is None:
+                # Should not happen because FunctionInfo always sets a single_fn
+                logger.warning("Skipping tool %s because single_fn is None", function_name)
+                continue
+
+            input_schema = tool_fn.input_schema
+            # Convert NoneType sentinel to None for FunctionGroup.add_function signature
+            if input_schema is type(None):  # noqa: E721
+                input_schema = None
+
+            # Add to group
+            logger.info("Adding tool %s to group", function_name)
+            group.add_function(
+                name=function_name,
+                description=description,
+                fn=single_fn,
+                input_schema=input_schema,
+                converters=tool_fn.converters,
+            )
+
+        yield group