datarobot-genai 0.1.64__tar.gz → 0.1.71__tar.gz

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.64
+Version: 0.1.71
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -32,7 +32,7 @@ Requires-Dist: aiohttp<4.0.0,>=3.9.0; extra == 'drmcp'
 Requires-Dist: aiosignal<2.0.0,>=1.3.1; extra == 'drmcp'
 Requires-Dist: boto3<2.0.0,>=1.34.0; extra == 'drmcp'
 Requires-Dist: datarobot-asgi-middleware<1.0.0,>=0.2.0; extra == 'drmcp'
-Requires-Dist: fastmcp==2.13.0.2; extra == 'drmcp'
+Requires-Dist: fastmcp<3.0.0,>=2.13.0.2; extra == 'drmcp'
 Requires-Dist: httpx<1.0.0,>=0.28.1; extra == 'drmcp'
 Requires-Dist: opentelemetry-api<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: opentelemetry-exporter-otlp-proto-http<2.0.0,>=1.22.0; extra == 'drmcp'

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.64"
+version = "0.1.71"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -84,7 +84,7 @@ drmcp = [
   "aiohttp>=3.9.0,<4.0.0",
   "aiohttp-retry>=2.8.3,<3.0.0",
   "aiosignal>=1.3.1,<2.0.0",
-  "fastmcp==2.13.0.2",
+  "fastmcp>=2.13.0.2,<3.0.0",
 ]
 
 [tool.hatch.build.targets.wheel]

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/core/agents/base.py

@@ -52,6 +52,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
         verbose: bool | str | None = True,
         timeout: int | None = 90,
         authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
         **_: Any,
     ) -> None:
         self.api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
@@ -68,6 +69,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
             self.verbose = bool(verbose)
         self._mcp_tools: list[TTool] = []
         self._authorization_context = authorization_context or {}
+        self._forwarded_headers: dict[str, str] = forwarded_headers or {}
 
     def set_mcp_tools(self, tools: list[TTool]) -> None:
         self._mcp_tools = tools
@@ -86,6 +88,11 @@ class BaseAgent(Generic[TTool], abc.ABC):
         """Return the authorization context for this agent."""
         return self._authorization_context
 
+    @property
+    def forwarded_headers(self) -> dict[str, str]:
+        """Return the forwarded headers for this agent."""
+        return self._forwarded_headers
+
     def litellm_api_base(self, deployment_id: str | None) -> str:
         return get_api_base(self.api_base, deployment_id)
 

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/core/custom_model.py

@@ -139,6 +139,11 @@ def chat_entrypoint(
     completion_create_params["authorization_context"] = resolve_authorization_context(
         completion_create_params, **kwargs
     )
+    # Keep only allowed headers from the forwarded_headers.
+    incoming_headers = kwargs.get("headers", {}) or {}
+    allowed_headers = {"x-datarobot-api-token", "x-datarobot-api-key"}
+    forwarded_headers = {k: v for k, v in incoming_headers.items() if k.lower() in allowed_headers}
+    completion_create_params["forwarded_headers"] = forwarded_headers
 
     # Instantiate user agent with all supplied completion params including auth context
     agent = agent_cls(**completion_create_params)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/core/mcp/common.py

@@ -13,15 +13,20 @@
 # limitations under the License.
 
 import json
+import logging
 import re
+from http import HTTPStatus
 from typing import Any
 from typing import Literal
 
+import requests
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
 from pydantic import field_validator
 
 from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
 
+logger = logging.getLogger(__name__)
+
 
 class MCPConfig(DataRobotAppFrameworkBaseSettings):
     """Configuration for MCP server connection.
@@ -37,6 +42,8 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
     datarobot_endpoint: str | None = None
     datarobot_api_token: str | None = None
     authorization_context: dict[str, Any] | None = None
+    forwarded_headers: dict[str, str] | None = None
+    mcp_server_port: int | None = None
 
     _auth_context_handler: AuthContextHeaderHandler | None = None
     _server_config: dict[str, Any] | None = None
@@ -47,17 +54,14 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
         if value is None:
             return None
 
-        if not isinstance(value, str):
-            msg = "external_mcp_headers must be a JSON string"
-            raise TypeError(msg)
-
         candidate = value.strip()
 
         try:
             json.loads(candidate)
-        except json.JSONDecodeError as exc:
+        except json.JSONDecodeError:
             msg = "external_mcp_headers must be valid JSON"
-            raise ValueError(msg) from exc
+            logger.warning(msg)
+            return None
 
         return candidate
 
@@ -67,15 +71,12 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
         if value is None:
             return None
 
-        if not isinstance(value, str):
-            msg = "mcp_deployment_id must be a string"
-            raise TypeError(msg)
-
         candidate = value.strip()
 
         if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
             msg = "mcp_deployment_id must be a valid 24-character hex ID"
-            raise ValueError(msg)
+            logger.warning(msg)
+            return None
 
         return candidate
 
@@ -110,6 +111,45 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
             # Authorization context not available (e.g., in tests)
             return {}
 
+    def _build_authenticated_headers(self) -> dict[str, str]:
+        """Build headers for authenticated requests.
+
+        Returns
+        -------
+            Dictionary containing forwarded headers (if available) and authentication headers.
+        """
+        headers: dict[str, str] = {}
+        if self.forwarded_headers:
+            headers.update(self.forwarded_headers)
+        headers.update(self._authorization_bearer_header())
+        headers.update(self._authorization_context_header())
+        return headers
+
+    def _check_localhost_server(self, url: str, timeout: float = 2.0) -> bool:
+        """Check if MCP server is running on localhost.
+
+        Parameters
+        ----------
+        url : str
+            The URL to check.
+        timeout : float, optional
+            Request timeout in seconds (default: 2.0).
+
+        Returns
+        -------
+        bool
+            True if server is running and responding with OK status, False otherwise.
+        """
+        try:
+            response = requests.get(url, timeout=timeout)
+            return (
+                response.status_code == HTTPStatus.OK
+                and response.json().get("message") == "DataRobot MCP Server is running"
+            )
+        except requests.RequestException as e:
+            logger.debug(f"Failed to connect to MCP server at {url}: {e}")
+            return False
+
     def _build_server_config(self) -> dict[str, Any] | None:
         """
         Get MCP server configuration.
@@ -119,20 +159,7 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
             Server configuration dict with url, transport, and optional headers,
             or None if not configured.
         """
-        if self.external_mcp_url:
-            # External MCP URL - no authentication needed
-            if self.external_mcp_headers:
-                headers = json.loads(self.external_mcp_headers)
-            else:
-                headers = {}
-
-            config = {
-                "url": self.external_mcp_url.rstrip("/"),
-                "transport": self.external_mcp_transport,
-                "headers": headers,
-            }
-            return config
-        elif self.mcp_deployment_id:
+        if self.mcp_deployment_id:
             # DataRobot deployment ID - requires authentication
             if self.datarobot_endpoint is None:
                 raise ValueError(
@@ -142,15 +169,15 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
                 raise ValueError(
                     "When using a DataRobot hosted MCP deployment, datarobot_api_token must be set."
                 )
+
             base_url = self.datarobot_endpoint.rstrip("/")
             if not base_url.endswith("/api/v2"):
-                base_url = base_url + "/api/v2"
+                base_url = f"{base_url}/api/v2"
+
             url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+            headers = self._build_authenticated_headers()
 
-            headers = {
-                **self._authorization_bearer_header(),
-                **self._authorization_context_header(),
-            }
+            logger.info(f"Using DataRobot hosted MCP deployment: {url}")
 
             return {
                 "url": url,
@@ -158,4 +185,34 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
                 "headers": headers,
             }
 
+        if self.external_mcp_url:
+            # External MCP URL - no authentication needed
+            headers = {}
+
+            # Merge external headers if provided
+            if self.external_mcp_headers:
+                external_headers = json.loads(self.external_mcp_headers)
+                headers.update(external_headers)
+
+            logger.info(f"Using external MCP URL: {self.external_mcp_url}")
+
+            return {
+                "url": self.external_mcp_url.rstrip("/"),
+                "transport": self.external_mcp_transport,
+                "headers": headers,
+            }
+
+        # No MCP configuration found, setup localhost if running locally
+        if self.mcp_server_port:
+            url = f"http://localhost:{self.mcp_server_port}"
+            if self._check_localhost_server(url):
+                headers = self._build_authenticated_headers()
+                logger.info(f"Using localhost MCP server: {url}")
+                return {
+                    "url": f"{url}/mcp",
+                    "transport": "streamable-http",
+                    "headers": headers,
+                }
+            logger.warning(f"MCP server is not running or not responding at {url}")
+
         return None

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/crewai/base.py

@@ -80,6 +80,37 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
         """
         raise NotImplementedError
 
+    def _extract_pipeline_interactions(self) -> MultiTurnSample | None:
+        """Extract pipeline interactions from event listener if available."""
+        if not hasattr(self, "event_listener"):
+            return None
+        try:
+            listener = getattr(self, "event_listener", None)
+            messages = getattr(listener, "messages", None) if listener is not None else None
+            return create_pipeline_interactions_from_messages(messages)
+        except Exception:
+            return None
+
+    def _extract_usage_metrics(self, crew_output: Any) -> UsageMetrics:
+        """Extract usage metrics from crew output."""
+        token_usage = getattr(crew_output, "token_usage", None)
+        if token_usage is not None:
+            return {
+                "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
+                "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
+                "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
+            }
+        return default_usage_metrics()
+
+    def _process_crew_output(
+        self, crew_output: Any
+    ) -> tuple[str, MultiTurnSample | None, UsageMetrics]:
+        """Process crew output into response tuple."""
+        response_text = str(crew_output.raw)
+        pipeline_interactions = self._extract_pipeline_interactions()
+        usage_metrics = self._extract_usage_metrics(crew_output)
+        return response_text, pipeline_interactions, usage_metrics
+
     async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
         """Run the CrewAI workflow with the provided completion parameters."""
         user_prompt_content = extract_user_prompt_content(completion_create_params)
@@ -93,6 +124,7 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
         # Use MCP context manager to handle connection lifecycle
         with mcp_tools_context(
             authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
         ) as mcp_tools:
             # Set MCP tools for all agents if MCP is not configured this is effectively a no-op
             self.set_mcp_tools(mcp_tools)
@@ -115,64 +147,13 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
                 async def _gen() -> AsyncGenerator[
                     tuple[str, MultiTurnSample | None, UsageMetrics]
                 ]:
-                    # Run kickoff in a worker thread.
                     crew_output = await asyncio.to_thread(
                         crew.kickoff,
                         inputs=self.make_kickoff_inputs(user_prompt_content),
                     )
-
-                    pipeline_interactions = None
-                    if hasattr(self, "event_listener"):
-                        try:
-                            listener = getattr(self, "event_listener", None)
-                            messages = (
-                                getattr(listener, "messages", None)
-                                if listener is not None
-                                else None
-                            )
-                            pipeline_interactions = create_pipeline_interactions_from_messages(
-                                messages
-                            )
-                        except Exception:
-                            pipeline_interactions = None
-
-                    token_usage = getattr(crew_output, "token_usage", None)
-                    if token_usage is not None:
-                        usage_metrics: UsageMetrics = {
-                            "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
-                            "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
-                            "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
-                        }
-                    else:
-                        usage_metrics = default_usage_metrics()
-
-                    # Finalize stream with empty chunk carrying interactions and usage
-                    yield "", pipeline_interactions, usage_metrics
+                    yield self._process_crew_output(crew_output)
 
                 return _gen()
 
-            # Non-streaming: run to completion and return final result
             crew_output = crew.kickoff(inputs=self.make_kickoff_inputs(user_prompt_content))
-
-            response_text = str(crew_output.raw)
-
-            pipeline_interactions = None
-            if hasattr(self, "event_listener"):
-                try:
-                    listener = getattr(self, "event_listener", None)
-                    messages = getattr(listener, "messages", None) if listener is not None else None
-                    pipeline_interactions = create_pipeline_interactions_from_messages(messages)
-                except Exception:
-                    pipeline_interactions = None
-
-            token_usage = getattr(crew_output, "token_usage", None)
-            if token_usage is not None:
-                usage_metrics: UsageMetrics = {
-                    "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
-                    "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
-                    "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
-                }
-            else:
-                usage_metrics = default_usage_metrics()
-
-            return response_text, pipeline_interactions, usage_metrics
+            return self._process_crew_output(crew_output)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/crewai/mcp.py

@@ -30,9 +30,13 @@ from datarobot_genai.core.mcp.common import MCPConfig
 @contextmanager
 def mcp_tools_context(
     authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
 ) -> Generator[list[Any], None, None]:
     """Context manager for MCP tools that handles connection lifecycle."""
-    config = MCPConfig(authorization_context=authorization_context)
+    config = MCPConfig(
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
+    )
     # If no MCP server configured, return empty tools list
     if not config.server_config:
         print("No MCP server configured, using empty tools list", flush=True)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -184,13 +184,17 @@ class DataRobotMCPServer:
             prompts = asyncio.run(self._mcp._list_prompts_mcp())
             resources = asyncio.run(self._mcp._list_resources_mcp())
 
-            self._logger.info(f"Registered tools: {len(tools)}")
+            tools_count = len(tools)
+            prompts_count = len(prompts)
+            resources_count = len(resources)
+
+            self._logger.info(f"Registered tools: {tools_count}")
             for tool in tools:
                 self._logger.info(f" > {tool.name}")
-            self._logger.info(f"Registered prompts: {len(prompts)}")
+            self._logger.info(f"Registered prompts: {prompts_count}")
             for prompt in prompts:
                 self._logger.info(f" > {prompt.name}")
-            self._logger.info(f"Registered resources: {len(resources)}")
+            self._logger.info(f"Registered resources: {resources_count}")
             for resource in resources:
                 self._logger.info(f" > {resource.name}")
 
@@ -209,6 +213,9 @@ class DataRobotMCPServer:
                         self._mcp,
                         self._mcp_transport,
                         port=self._config.mcp_server_port,
+                        tools_count=tools_count,
+                        prompts_count=prompts_count,
+                        resources_count=resources_count,
                     )
 
                 if self._mcp_transport == "stdio":

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py

@@ -38,7 +38,7 @@ def _apply_green(text: str) -> str:
     return "\n".join(colored_lines)
 
 
-DR_LOGO_ASCII = _apply_green("""\
+DR_LOGO_ASCII = _apply_green(r"""
  ____        _        ____       _           _   
 |  _ \  __ _| |_ __ _|  _ \ ___ | |__   ___ | |_ 
 | | | |/ _` | __/ _` | |_) / _ \| '_ \ / _ \| __|
@@ -54,6 +54,9 @@ def log_server_custom_banner(
     host: str | None = None,
     port: int | None = None,
     path: str | None = None,
+    tools_count: int | None = None,
+    prompts_count: int | None = None,
+    resources_count: int | None = None,
 ) -> None:
     """
     Create and log a formatted banner with server information and logo.
@@ -64,13 +67,20 @@ def log_server_custom_banner(
         host: Host address (for HTTP transports)
         port: Port number (for HTTP transports)
         path: Server path (for HTTP transports)
+        tools_count: Number of tools registered
+        prompts_count: Number of prompts registered
+        resources_count: Number of resources registered
     """
     # Create the logo text
     # Use Text with no_wrap and markup disabled to preserve ANSI escape codes
     logo_text = Text.from_ansi(DR_LOGO_ASCII, no_wrap=True)
 
     # Create the main title
-    title_text = Text(f"DataRobot MCP Server {datarobot_genai_version}", style="bold green")
+    title_text = Text(f"DataRobot MCP Server {datarobot_genai_version}", style="dim green")
+    stats_text = Text(
+        f"{tools_count} tools, {prompts_count} prompts, {resources_count} resources",
+        style="bold green",
+    )
 
     # Create the information table
     info_table = Table.grid(padding=(0, 1))
@@ -107,6 +117,7 @@ def log_server_custom_banner(
         Align.center(logo_text),
         "",
         Align.center(title_text),
+        Align.center(stats_text),
         "",
         "",
         Align.center(info_table),

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py

@@ -18,6 +18,8 @@ from fastmcp.prompts.prompt import Prompt
 
 from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template
 from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_version
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_versions
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_templates
 from datarobot_genai.drmcp.core.dynamic_prompts.register import (
     register_prompt_from_datarobot_prompt_management,
 )
@@ -83,3 +85,46 @@ async def delete_registered_prompt_template(prompt_template_id: str) -> bool:
         f"version {prompt_template_version_id}"
     )
     return True
+
+
+async def refresh_registered_prompt_template() -> None:
+    """Refresh all registered prompt templates in the MCP instance."""
+    prompt_templates = get_datarobot_prompt_templates()
+    prompt_templates_ids = {p.id for p in prompt_templates}
+    prompt_templates_versions = get_datarobot_prompt_template_versions(list(prompt_templates_ids))
+
+    mcp_prompt_templates_mappings = await mcp.get_prompt_mapping()
+
+    for prompt_template in prompt_templates:
+        prompt_template_versions = prompt_templates_versions.get(prompt_template.id)
+        if not prompt_template_versions:
+            continue
+
+        latest_version = max(prompt_template_versions, key=lambda v: v.version)
+
+        if prompt_template.id not in mcp_prompt_templates_mappings:
+            # New prompt template -> add
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        mcp_prompt_template_version, mcp_prompt = mcp_prompt_templates_mappings[prompt_template.id]
+
+        if mcp_prompt_template_version != latest_version:
+            # Current version saved in MCP is not the latest one => update it
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        # Else => mcp_prompt_template_version == latest_version
+        # For now it means nothing changed as there's no possibility to edit promp template version.
+
+    for mcp_prompt_template_id, (
+        mcp_prompt_template_version_id,
+        _,
+    ) in mcp_prompt_templates_mappings.items():
+        if mcp_prompt_template_id not in prompt_templates_ids:
+            # We need to also delete prompt templates that are
+            await mcp.remove_prompt_mapping(mcp_prompt_template_id, mcp_prompt_template_version_id)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/drmcp/core/mcp_server_tools.py

@@ -51,7 +51,7 @@ async def list_tools_by_tags(tags: list[str] | None = None, match_all: bool = Fa
     -------
         A formatted string listing tools that match the tag criteria.
     """
-    tools = await mcp._list_tools_mcp(tags=tags, match_all=match_all)
+    tools = await mcp.list_tools(tags=tags, match_all=match_all)
 
     if not tools:
         if tags:
@@ -95,7 +95,7 @@ async def get_tool_info_by_name(tool_name: str) -> str:
     -------
         A formatted string with detailed information about the tool.
     """
-    all_tools = await mcp._list_tools_mcp()
+    all_tools = await mcp.list_tools()
 
     for tool in all_tools:
         if tool.name == tool_name:

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/drmcp/core/routes.py

@@ -19,6 +19,7 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse
 
 from .dynamic_prompts.controllers import delete_registered_prompt_template
+from .dynamic_prompts.controllers import refresh_registered_prompt_template
 from .dynamic_prompts.controllers import register_prompt_from_prompt_template_id_and_version
 from .dynamic_tools.deployment.controllers import delete_registered_tool_deployment
 from .dynamic_tools.deployment.controllers import get_registered_tool_deployments
@@ -418,6 +419,18 @@ def register_routes(mcp: TaggedFastMCP) -> None:
             )
         except Exception as e:
             return JSONResponse(
-                status_code=HTTPStatus.BAD_REQUEST,
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 content={"error": f"Failed to add prompt template: {str(e)}"},
             )
+
+    @mcp.custom_route(prefix_mount_path("/registeredPrompts"), methods=["PUT"])
+    async def refresh_prompt_templates(_: Request) -> JSONResponse:
+        """Refresh prompt templates."""
+        try:
+            await refresh_registered_prompt_template()
+            return JSONResponse(status_code=HTTPStatus.NO_CONTENT, content=None)
+        except Exception as e:
+            return JSONResponse(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                content={"error": f"Failed to refresh prompt templates: {str(e)}"},
+            )

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/langgraph/agent.py

@@ -84,26 +84,45 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
             async def wrapped_generator() -> AsyncGenerator[
                 tuple[str, Any | None, UsageMetrics], None
             ]:
-                async with mcp_tools_context(
-                    authorization_context=self._authorization_context,
-                ) as mcp_tools:
-                    self.set_mcp_tools(mcp_tools)
-                    result = await self._invoke(completion_create_params)
-
-                    # Yield all items from the result generator
-                    # The context will be closed when this generator is exhausted
-                    # Cast to async generator since we know stream=True means it's a generator
-                    result_generator = cast(
-                        AsyncGenerator[tuple[str, Any | None, UsageMetrics], None], result
-                    )
-                    async for item in result_generator:
-                        yield item
+                try:
+                    async with mcp_tools_context(
+                        authorization_context=self._authorization_context,
+                        forwarded_headers=self.forwarded_headers,
+                    ) as mcp_tools:
+                        self.set_mcp_tools(mcp_tools)
+                        result = await self._invoke(completion_create_params)
+
+                        # Yield all items from the result generator
+                        # The context will be closed when this generator is exhausted
+                        # Cast to async generator since we know stream=True means it's a generator
+                        result_generator = cast(
+                            AsyncGenerator[tuple[str, Any | None, UsageMetrics], None], result
+                        )
+                        async for item in result_generator:
+                            yield item
+                except RuntimeError as e:
+                    error_message = str(e).lower()
+                    if "different task" in error_message and "cancel scope" in error_message:
+                        # Due to anyio task group constraints when consuming async generators
+                        # across task boundaries, we cannot always clean up properly.
+                        # The underlying HTTP client/connection pool should handle resource cleanup
+                        # via timeouts and connection pooling, but this
+                        # may lead to delayed resource release.
+                        logger.debug(
+                            "MCP context cleanup attempted in different task. "
+                            "This is a limitation when consuming async generators "
+                            "across task boundaries."
+                        )
+                    else:
+                        # Re-raise if it's a different RuntimeError
+                        raise
 
             return wrapped_generator()
         else:
             # For non-streaming, use async with directly
             async with mcp_tools_context(
                 authorization_context=self._authorization_context,
+                forwarded_headers=self.forwarded_headers,
             ) as mcp_tools:
                 self.set_mcp_tools(mcp_tools)
                 result = await self._invoke(completion_create_params)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/langgraph/mcp.py

@@ -28,6 +28,7 @@ from datarobot_genai.core.mcp.common import MCPConfig
 @asynccontextmanager
 async def mcp_tools_context(
     authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
 ) -> AsyncGenerator[list[BaseTool], None]:
     """Yield a list of LangChain BaseTool instances loaded via MCP.
 
@@ -37,8 +38,13 @@ async def mcp_tools_context(
     ----------
     authorization_context : dict[str, Any] | None
         Authorization context to use for MCP connections
+    forwarded_headers : dict[str, str] | None
+        Forwarded headers, e.g. x-datarobot-api-key to use for MCP authentication
     """
-    mcp_config = MCPConfig(authorization_context=authorization_context)
+    mcp_config = MCPConfig(
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
+    )
     server_config = mcp_config.server_config
 
     if not server_config:

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/llama_index/base.py

@@ -84,6 +84,7 @@ class LlamaIndexAgent(BaseAgent[BaseTool], abc.ABC):
         # Load MCP tools (if configured) asynchronously before building workflow
         mcp_tools = await load_mcp_tools(
             authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
         )
         self.set_mcp_tools(mcp_tools)
 

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/llama_index/mcp.py

@@ -30,18 +30,23 @@ from datarobot_genai.core.mcp.common import MCPConfig
 
 async def load_mcp_tools(
     authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
 ) -> list[Any]:
     """
     Asynchronously load MCP tools for LlamaIndex.
 
     Args:
         authorization_context: Optional authorization context for MCP connections
+        forwarded_headers: Optional forwarded headers, e.g. x-datarobot-api-key for MCP auth
 
     Returns
     -------
         List of MCP tools, or empty list if no MCP configuration is present.
     """
-    config = MCPConfig(authorization_context=authorization_context)
+    config = MCPConfig(
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
+    )
     server_params = config.server_config
 
     if not server_params:

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/nat/datarobot_llm_clients.py

@@ -23,6 +23,7 @@ from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
 
+from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMDeploymentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMGatewayModelConfig
 from ..nat.datarobot_llm_providers import DataRobotNIMModelConfig
@@ -75,6 +76,7 @@ async def datarobot_llm_gateway_langchain(
     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/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -85,7 +87,8 @@ async def datarobot_llm_gateway_crewai(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["base_url"] = config["base_url"].removesuffix("/api/v2")
     yield LLM(**config)
 
@@ -97,7 +100,8 @@ async def datarobot_llm_gateway_llamaindex(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
     yield DataRobotLiteLLM(**config)
 
@@ -109,11 +113,12 @@ async def datarobot_llm_deployment_langchain(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
     config = llm_config.model_dump(
-        exclude={"type", "thinking", "datarobot_endpoint", "llm_deployment_id"},
+        exclude={"type", "thinking"},
         by_alias=True,
         exclude_none=True,
     )
     config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -128,7 +133,8 @@ async def datarobot_llm_deployment_crewai(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield LLM(**config)
 
@@ -144,7 +150,8 @@ async def datarobot_llm_deployment_llamaindex(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield DataRobotLiteLLM(**config)
 
@@ -159,6 +166,7 @@ async def datarobot_nim_langchain(
         exclude_none=True,
     )
     config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -171,7 +179,8 @@ async def datarobot_nim_crewai(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield LLM(**config)
 
@@ -185,6 +194,56 @@ async def datarobot_nim_llamaindex(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield DataRobotLiteLLM(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN
+)
+async def datarobot_llm_component_langchain(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    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")
+    yield DataRobotChatOpenAI(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI
+)
+async def datarobot_llm_component_crewai(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["base_url"] = config["base_url"].removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    yield LLM(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX
+)
+async def datarobot_llm_component_llamaindex(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    yield DataRobotLiteLLM(**config)

{datarobot_genai-0.1.64 → datarobot_genai-0.1.71}/src/datarobot_genai/nat/datarobot_llm_providers.py

@@ -32,11 +32,43 @@ class Config(DataRobotAppFrameworkBaseSettings):
     datarobot_api_token: str | None = None
     llm_deployment_id: str | None = None
     nim_deployment_id: str | None = None
+    use_datarobot_llm_gateway: bool = False
+    llm_default_model: str | None = None
 
 
 config = Config()
 
 
+class DataRobotLLMComponentModelConfig(OpenAIModelConfig, name="datarobot-llm-component"):  # type: ignore[call-arg]
+    """A DataRobot LLM provider to be used with an LLM client."""
+
+    api_key: str | None = Field(
+        default=config.datarobot_api_token, description="DataRobot API key."
+    )
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint.rstrip("/")
+        if config.use_datarobot_llm_gateway
+        else config.datarobot_endpoint + f"/deployments/{config.llm_deployment_id}",
+        description="DataRobot LLM URL.",
+    )
+    model_name: str = Field(
+        validation_alias=AliasChoices("model_name", "model"),
+        serialization_alias="model",
+        description="The model name.",
+        default=config.llm_default_model or "datarobot-deployed-llm",
+    )
+    use_datarobot_llm_gateway: bool = config.use_datarobot_llm_gateway
+
+
+@register_llm_provider(config_type=DataRobotLLMComponentModelConfig)
+async def datarobot_llm_component(
+    config: DataRobotLLMComponentModelConfig, _builder: Builder
+) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot LLM Component for use with an LLM client."
+    )
+
+
 class DataRobotLLMGatewayModelConfig(OpenAIModelConfig, name="datarobot-llm-gateway"):  # type: ignore[call-arg]
     """A DataRobot LLM provider to be used with an LLM client."""