datarobot-genai 0.1.66__tar.gz → 0.1.67__tar.gz

{datarobot_genai-0.1.66 → datarobot_genai-0.1.67}/PKG-INFO

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

{datarobot_genai-0.1.66 → datarobot_genai-0.1.67}/pyproject.toml

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

{datarobot_genai-0.1.66 → datarobot_genai-0.1.67}/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.66 → datarobot_genai-0.1.67}/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.66 → datarobot_genai-0.1.67}/src/datarobot_genai/core/mcp/common.py

@@ -16,6 +16,7 @@ import json
 import re
 from typing import Any
 from typing import Literal
+from urllib.parse import urlparse
 
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
 from pydantic import field_validator
@@ -37,6 +38,7 @@ 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
 
     _auth_context_handler: AuthContextHeaderHandler | None = None
     _server_config: dict[str, Any] | None = None
@@ -121,18 +123,33 @@ class MCPConfig(DataRobotAppFrameworkBaseSettings):
         """
         if self.external_mcp_url:
             # External MCP URL - no authentication needed
+            headers: dict[str, str] = {}
+
+            # Forward headers for localhost connections
+            if self.forwarded_headers:
+                try:
+                    parsed_url = urlparse(self.external_mcp_url)
+                    hostname = parsed_url.hostname or ""
+                    # Check if hostname is localhost or 127.0.0.1
+                    if hostname in ("localhost", "127.0.0.1", "::1"):
+                        headers.update(self.forwarded_headers)
+                except Exception:
+                    # If URL parsing fails, fall back to simple string check
+                    if "localhost" in self.external_mcp_url or "127.0.0.1" in self.external_mcp_url:
+                        headers.update(self.forwarded_headers)
+
+            # Merge external headers if provided
             if self.external_mcp_headers:
-                headers = json.loads(self.external_mcp_headers)
-            else:
-                headers = {}
+                external_headers = json.loads(self.external_mcp_headers)
+                headers.update(external_headers)
 
-            config = {
+            return {
                 "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 +159,21 @@ 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._authorization_bearer_header(),
-                **self._authorization_context_header(),
-            }
+            # Start with forwarded headers if available
+            headers = {}
+            if self.forwarded_headers:
+                headers.update(self.forwarded_headers)
+
+            # Add authentication headers
+            headers.update(self._authorization_bearer_header())
+            headers.update(self._authorization_context_header())
 
             return {
                 "url": url,

{datarobot_genai-0.1.66 → datarobot_genai-0.1.67}/src/datarobot_genai/crewai/base.py

@@ -93,6 +93,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)

{datarobot_genai-0.1.66 → datarobot_genai-0.1.67}/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.66 → datarobot_genai-0.1.67}/src/datarobot_genai/langgraph/agent.py

@@ -86,6 +86,7 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
             ]:
                 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)
@@ -104,6 +105,7 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
             # 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.66 → datarobot_genai-0.1.67}/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.66 → datarobot_genai-0.1.67}/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.66 → datarobot_genai-0.1.67}/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: