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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.55
+Version: 0.1.64
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -8,65 +8,65 @@ License: Apache-2.0
 License-File: AUTHORS
 License-File: LICENSE
 Requires-Python: <3.13,>=3.10
-Requires-Dist: datarobot-drum>=1.17.5
-Requires-Dist: datarobot-predict>=1.13.2
-Requires-Dist: datarobot>=3.9.1
-Requires-Dist: openai>=1.76.2
-Requires-Dist: opentelemetry-instrumentation-aiohttp-client>=0.43b0
-Requires-Dist: opentelemetry-instrumentation-httpx>=0.43b0
-Requires-Dist: opentelemetry-instrumentation-openai>=0.40.5
-Requires-Dist: opentelemetry-instrumentation-requests>=0.43b0
-Requires-Dist: pandas>=2.2.3
-Requires-Dist: pyjwt>=2.10.1
-Requires-Dist: pypdf>=6.1.3
-Requires-Dist: ragas>=0.3.8
-Requires-Dist: requests>=2.32.4
+Requires-Dist: datarobot-drum<2.0.0,>=1.17.5
+Requires-Dist: datarobot-predict<2.0.0,>=1.13.2
+Requires-Dist: datarobot<4.0.0,>=3.9.1
+Requires-Dist: openai<2.0.0,>=1.76.2
+Requires-Dist: opentelemetry-instrumentation-aiohttp-client<1.0.0,>=0.43b0
+Requires-Dist: opentelemetry-instrumentation-httpx<1.0.0,>=0.43b0
+Requires-Dist: opentelemetry-instrumentation-openai<1.0.0,>=0.40.5
+Requires-Dist: opentelemetry-instrumentation-requests<1.0.0,>=0.43b0
+Requires-Dist: pandas<3.0.0,>=2.2.3
+Requires-Dist: pyjwt<3.0.0,>=2.10.1
+Requires-Dist: pypdf<7.0.0,>=6.1.3
+Requires-Dist: ragas<0.4.0,>=0.3.8
+Requires-Dist: requests<3.0.0,>=2.32.4
 Provides-Extra: crewai
-Requires-Dist: crewai-tools[mcp]>=0.69.0; extra == 'crewai'
-Requires-Dist: crewai>=0.193.2; extra == 'crewai'
-Requires-Dist: opentelemetry-instrumentation-crewai>=0.40.5; extra == 'crewai'
-Requires-Dist: pybase64>=1.4.2; extra == 'crewai'
+Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
+Requires-Dist: crewai<1.0.0,>=0.193.2; extra == 'crewai'
+Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'
+Requires-Dist: pybase64<2.0.0,>=1.4.2; extra == 'crewai'
 Provides-Extra: drmcp
-Requires-Dist: aiohttp-retry>=2.8.3; extra == 'drmcp'
-Requires-Dist: aiohttp>=3.9.0; extra == 'drmcp'
-Requires-Dist: aiosignal>=1.3.1; extra == 'drmcp'
-Requires-Dist: boto3>=1.34.0; extra == 'drmcp'
-Requires-Dist: datarobot-asgi-middleware>=0.2.0; extra == 'drmcp'
+Requires-Dist: aiohttp-retry<3.0.0,>=2.8.3; extra == 'drmcp'
+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: httpx>=0.28.1; extra == 'drmcp'
-Requires-Dist: opentelemetry-api>=1.22.0; extra == 'drmcp'
-Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.22.0; extra == 'drmcp'
-Requires-Dist: opentelemetry-exporter-otlp>=1.22.0; extra == 'drmcp'
-Requires-Dist: opentelemetry-sdk>=1.22.0; extra == 'drmcp'
-Requires-Dist: pydantic-settings>=2.1.0; extra == 'drmcp'
-Requires-Dist: pydantic>=2.6.1; extra == 'drmcp'
-Requires-Dist: python-dotenv>=1.1.0; 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'
+Requires-Dist: opentelemetry-exporter-otlp<2.0.0,>=1.22.0; extra == 'drmcp'
+Requires-Dist: opentelemetry-sdk<2.0.0,>=1.22.0; extra == 'drmcp'
+Requires-Dist: pydantic-settings<3.0.0,>=2.1.0; extra == 'drmcp'
+Requires-Dist: pydantic<3.0.0,>=2.6.1; extra == 'drmcp'
+Requires-Dist: python-dotenv<2.0.0,>=1.1.0; extra == 'drmcp'
 Provides-Extra: langgraph
-Requires-Dist: langchain-mcp-adapters>=0.1.12; extra == 'langgraph'
-Requires-Dist: langgraph-prebuilt>=0.2.3; extra == 'langgraph'
-Requires-Dist: langgraph>=0.4.10; extra == 'langgraph'
-Requires-Dist: opentelemetry-instrumentation-langchain>=0.40.5; extra == 'langgraph'
+Requires-Dist: langchain-mcp-adapters<0.2.0,>=0.1.12; extra == 'langgraph'
+Requires-Dist: langgraph-prebuilt<0.7.0,>=0.2.3; extra == 'langgraph'
+Requires-Dist: langgraph<0.7.0,>=0.4.10; extra == 'langgraph'
+Requires-Dist: opentelemetry-instrumentation-langchain<1.0.0,>=0.40.5; extra == 'langgraph'
 Provides-Extra: llamaindex
 Requires-Dist: llama-index-core<0.14.0,>=0.13.6; extra == 'llamaindex'
-Requires-Dist: llama-index-llms-langchain>=0.6.1; extra == 'llamaindex'
-Requires-Dist: llama-index-llms-litellm>=0.4.1; extra == 'llamaindex'
-Requires-Dist: llama-index-llms-openai>=0.3.38; extra == 'llamaindex'
-Requires-Dist: llama-index-tools-mcp>=0.1.0; extra == 'llamaindex'
+Requires-Dist: llama-index-llms-langchain<0.8.0,>=0.6.1; extra == 'llamaindex'
+Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'llamaindex'
+Requires-Dist: llama-index-llms-openai<0.6.0,>=0.3.38; extra == 'llamaindex'
+Requires-Dist: llama-index-tools-mcp<0.5.0,>=0.1.0; extra == 'llamaindex'
 Requires-Dist: llama-index<0.14.0,>=0.13.6; extra == 'llamaindex'
-Requires-Dist: opentelemetry-instrumentation-llamaindex>=0.40.5; extra == 'llamaindex'
-Requires-Dist: pypdf>=6.0.0; extra == 'llamaindex'
+Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'llamaindex'
+Requires-Dist: pypdf<7.0.0,>=6.0.0; extra == 'llamaindex'
 Provides-Extra: nat
-Requires-Dist: llama-index-llms-litellm>=0.4.1; extra == 'nat'
+Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
 Requires-Dist: nvidia-nat-crewai==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'
-Requires-Dist: opentelemetry-instrumentation-crewai>=0.40.5; extra == 'nat'
-Requires-Dist: opentelemetry-instrumentation-langchain>=0.40.5; extra == 'nat'
-Requires-Dist: opentelemetry-instrumentation-llamaindex>=0.40.5; extra == 'nat'
+Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'nat'
+Requires-Dist: opentelemetry-instrumentation-langchain<1.0.0,>=0.40.5; extra == 'nat'
+Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'nat'
 Provides-Extra: pydanticai
-Requires-Dist: logfire>=4.6.0; extra == 'pydanticai'
-Requires-Dist: pydantic-ai>=1.0.5; extra == 'pydanticai'
+Requires-Dist: logfire<5.0.0,>=4.6.0; extra == 'pydanticai'
+Requires-Dist: pydantic-ai<1.9.0,>=1.0.5; extra == 'pydanticai'
 Description-Content-Type: text/markdown
 
 <p align="center">

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

@@ -4,26 +4,26 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.55"
+version = "0.1.64"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
 license = { text = "Apache-2.0" }
 authors = [{ name = "DataRobot, Inc." }]
 dependencies = [
-  "requests>=2.32.4",
-  "datarobot>=3.9.1",
-  "datarobot-predict>=1.13.2",
-  "openai>=1.76.2",
-  "pandas>=2.2.3",
-  "ragas>=0.3.8",
-  "pyjwt>=2.10.1",
-  "pypdf>=6.1.3",  # CVE BUZZOK-28182
-  "datarobot-drum>=1.17.5",
-  "opentelemetry-instrumentation-requests>=0.43b0",
-  "opentelemetry-instrumentation-aiohttp-client>=0.43b0",
-  "opentelemetry-instrumentation-httpx>=0.43b0",
-  "opentelemetry-instrumentation-openai>=0.40.5",
+  "requests>=2.32.4,<3.0.0",
+  "datarobot>=3.9.1,<4.0.0",
+  "datarobot-predict>=1.13.2,<2.0.0",
+  "openai>=1.76.2,<2.0.0",
+  "pandas>=2.2.3,<3.0.0",
+  "ragas>=0.3.8,<0.4.0",
+  "pyjwt>=2.10.1,<3.0.0",
+  "pypdf>=6.1.3,<7.0.0",  # CVE BUZZOK-28182
+  "datarobot-drum>=1.17.5,<2.0.0",
+  "opentelemetry-instrumentation-requests>=0.43b0,<1.0.0",
+  "opentelemetry-instrumentation-aiohttp-client>=0.43b0,<1.0.0",
+  "opentelemetry-instrumentation-httpx>=0.43b0,<1.0.0",
+  "opentelemetry-instrumentation-openai>=0.40.5,<1.0.0",
 ]
 
 [project.urls]
@@ -35,55 +35,55 @@ datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
 
 [project.optional-dependencies]
 crewai = [
-  "crewai>=0.193.2",
-  "crewai-tools[mcp]>=0.69.0",
-  "opentelemetry-instrumentation-crewai>=0.40.5",
-  "pybase64>=1.4.2",
+  "crewai>=0.193.2,<1.0.0",
+  "crewai-tools[mcp]>=0.69.0,<0.77.0",
+  "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
+  "pybase64>=1.4.2,<2.0.0",
 ]
 langgraph = [
-  "langchain-mcp-adapters>=0.1.12",
-  "langgraph>=0.4.10",
-  "langgraph-prebuilt>=0.2.3",
-  "opentelemetry-instrumentation-langchain>=0.40.5",
+  "langchain-mcp-adapters>=0.1.12,<0.2.0",
+  "langgraph>=0.4.10,<0.7.0",
+  "langgraph-prebuilt>=0.2.3,<0.7.0",
+  "opentelemetry-instrumentation-langchain>=0.40.5,<1.0.0",
 ]
 llamaindex = [
   "llama-index>=0.13.6,<0.14.0",
   "llama-index-core>=0.13.6,<0.14.0",
-  "llama-index-llms-langchain>=0.6.1",
-  "llama-index-llms-litellm>=0.4.1",  # Sync nat dependency if possible too
-  "llama-index-llms-openai>=0.3.38",
-  "llama-index-tools-mcp>=0.1.0",
-  "opentelemetry-instrumentation-llamaindex>=0.40.5",
-  "pypdf>=6.0.0",
+  "llama-index-llms-langchain>=0.6.1,<0.8.0",
+  "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Sync nat dependency if possible too
+  "llama-index-llms-openai>=0.3.38,<0.6.0",
+  "llama-index-tools-mcp>=0.1.0,<0.5.0",
+  "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",
+  "pypdf>=6.0.0,<7.0.0",
 ]
 nat = [
   "nvidia-nat==1.3.0; python_version >= '3.11'",
   "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
   "nvidia-nat-crewai==1.3.0; python_version >= '3.11'",
   "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
-  "llama-index-llms-litellm>=0.4.1",  # Need this to support datarobot-llm plugin
-  "opentelemetry-instrumentation-crewai>=0.40.5",
-  "opentelemetry-instrumentation-llamaindex>=0.40.5",
-  "opentelemetry-instrumentation-langchain>=0.40.5",
+  "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Need this to support datarobot-llm plugin
+  "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
+  "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",
+  "opentelemetry-instrumentation-langchain>=0.40.5,<1.0.0",
 ]
 pydanticai = [
-  "logfire>=4.6.0",
-  "pydantic-ai>=1.0.5",
+  "logfire>=4.6.0,<5.0.0",
+  "pydantic-ai>=1.0.5,<1.9.0",
 ]
 drmcp = [
-  "datarobot-asgi-middleware>=0.2.0",
-  "python-dotenv>=1.1.0",
-  "boto3>=1.34.0",
-  "httpx>=0.28.1",
-  "pydantic>=2.6.1",
-  "pydantic-settings>=2.1.0",
-  "opentelemetry-api>=1.22.0",
-  "opentelemetry-sdk>=1.22.0",
-  "opentelemetry-exporter-otlp>=1.22.0",
-  "opentelemetry-exporter-otlp-proto-http>=1.22.0",
-  "aiohttp>=3.9.0",
-  "aiohttp-retry>=2.8.3",
-  "aiosignal>=1.3.1",
+  "datarobot-asgi-middleware>=0.2.0,<1.0.0",
+  "python-dotenv>=1.1.0,<2.0.0",
+  "boto3>=1.34.0,<2.0.0",
+  "httpx>=0.28.1,<1.0.0",
+  "pydantic>=2.6.1,<3.0.0",
+  "pydantic-settings>=2.1.0,<3.0.0",
+  "opentelemetry-api>=1.22.0,<2.0.0",
+  "opentelemetry-sdk>=1.22.0,<2.0.0",
+  "opentelemetry-exporter-otlp>=1.22.0,<2.0.0",
+  "opentelemetry-exporter-otlp-proto-http>=1.22.0,<2.0.0",
+  "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",
 ]
 

datarobot_genai-0.1.64/src/datarobot_genai/core/mcp/common.py

@@ -0,0 +1,161 @@
+# 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 json
+import re
+from typing import Any
+from typing import Literal
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from pydantic import field_validator
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+
+class MCPConfig(DataRobotAppFrameworkBaseSettings):
+    """Configuration for MCP server connection.
+
+    Derived values are exposed as properties rather than stored, avoiding
+    Pydantic field validation/serialization concerns for internal helpers.
+    """
+
+    external_mcp_url: str | None = None
+    external_mcp_headers: str | None = None
+    external_mcp_transport: Literal["sse", "streamable-http"] = "streamable-http"
+    mcp_deployment_id: str | None = None
+    datarobot_endpoint: str | None = None
+    datarobot_api_token: str | None = None
+    authorization_context: dict[str, Any] | None = None
+
+    _auth_context_handler: AuthContextHeaderHandler | None = None
+    _server_config: dict[str, Any] | None = None
+
+    @field_validator("external_mcp_headers", mode="before")
+    @classmethod
+    def validate_external_mcp_headers(cls, value: str | None) -> str | None:
+        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:
+            msg = "external_mcp_headers must be valid JSON"
+            raise ValueError(msg) from exc
+
+        return candidate
+
+    @field_validator("mcp_deployment_id", mode="before")
+    @classmethod
+    def validate_mcp_deployment_id(cls, value: str | None) -> str | None:
+        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)
+
+        return candidate
+
+    def _authorization_bearer_header(self) -> dict[str, str]:
+        """Return Authorization header with Bearer token or empty dict."""
+        if not self.datarobot_api_token:
+            return {}
+        auth = (
+            self.datarobot_api_token
+            if self.datarobot_api_token.startswith("Bearer ")
+            else f"Bearer {self.datarobot_api_token}"
+        )
+        return {"Authorization": auth}
+
+    @property
+    def auth_context_handler(self) -> AuthContextHeaderHandler:
+        if self._auth_context_handler is None:
+            self._auth_context_handler = AuthContextHeaderHandler()
+        return self._auth_context_handler
+
+    @property
+    def server_config(self) -> dict[str, Any] | None:
+        if self._server_config is None:
+            self._server_config = self._build_server_config()
+        return self._server_config
+
+    def _authorization_context_header(self) -> dict[str, str]:
+        """Return X-DataRobot-Authorization-Context header or empty dict."""
+        try:
+            return self.auth_context_handler.get_header(self.authorization_context)
+        except (LookupError, RuntimeError):
+            # Authorization context not available (e.g., in tests)
+            return {}
+
+    def _build_server_config(self) -> dict[str, Any] | None:
+        """
+        Get MCP server configuration.
+
+        Returns
+        -------
+            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:
+            # DataRobot deployment ID - requires authentication
+            if self.datarobot_endpoint is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_endpoint must be set."
+                )
+            if self.datarobot_api_token is None:
+                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"
+            url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+
+            headers = {
+                **self._authorization_bearer_header(),
+                **self._authorization_context_header(),
+            }
+
+            return {
+                "url": url,
+                "transport": "streamable-http",
+                "headers": headers,
+            }
+
+        return None

{datarobot_genai-0.1.55 → datarobot_genai-0.1.64}/src/datarobot_genai/core/telemetry_agent.py

@@ -28,10 +28,22 @@ logging.getLogger("opentelemetry.instrumentation.instrumentor").setLevel(logging
 logger = logging.getLogger(__name__)
 
 # Internal instrumentation state to avoid 'global' mutation warnings
-_INSTRUMENTATION_STATE = {"http": False, "openai": False}
+_INSTRUMENTATION_STATE = {"http": False, "openai": False, "threading": False}
 _INSTRUMENTED_FRAMEWORKS: set[str] = set()
 
 
+def _instrument_threading() -> None:
+    if _INSTRUMENTATION_STATE["threading"]:
+        return
+    try:
+        threading_module = importlib.import_module("opentelemetry.instrumentation.threading")
+        threading_instrumentor = getattr(threading_module, "ThreadingInstrumentor")
+        threading_instrumentor().instrument()
+        _INSTRUMENTATION_STATE["threading"] = True
+    except Exception as e:
+        logger.debug(f"threading instrumentation skipped: {e}")
+
+
 def _instrument_http_clients() -> None:
     if _INSTRUMENTATION_STATE["http"]:
         return
@@ -107,6 +119,7 @@ def instrument(
     os.environ.setdefault("RAGAS_DO_NOT_TRACK", "true")
     os.environ.setdefault("DEEPEVAL_TELEMETRY_OPT_OUT", "YES")
 
+    _instrument_threading()
     _instrument_http_clients()
     _instrument_openai()
     if framework:

{datarobot_genai-0.1.55 → datarobot_genai-0.1.64}/src/datarobot_genai/core/utils/auth.py

@@ -16,9 +16,14 @@ import warnings
 from typing import Any
 
 import jwt
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.identity import Identity
+from datarobot.auth.oauth import AsyncOAuthComponent
 from datarobot.auth.session import AuthCtx
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from datarobot.models.genai.agent.auth import ToolAuth
 from datarobot.models.genai.agent.auth import get_authorization_context
+from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +32,13 @@ class AuthContextConfig(DataRobotAppFrameworkBaseSettings):
     session_secret_key: str = ""
 
 
+class DRAppCtx(BaseModel):
+    """DataRobot application context from authorization metadata."""
+
+    email: str | None = None
+    api_key: str | None = None
+
+
 class AuthContextHeaderHandler:
     """Manages encoding and decoding of authorization context into JWT tokens.
 
@@ -146,6 +158,7 @@ class AuthContextHeaderHandler:
 
         auth_ctx_dict = self.decode(token)
         if not auth_ctx_dict:
+            logger.debug("Failed to decode auth context from token")
             return None
 
         try:
@@ -153,3 +166,54 @@ class AuthContextHeaderHandler:
         except Exception as e:
             logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
             return None
+
+
+class AsyncOAuthTokenProvider:
+    """Manages OAuth access tokens using generic OAuth client."""
+
+    def __init__(self, auth_ctx: AuthCtx) -> None:
+        self.auth_ctx = auth_ctx
+        self.oauth_client = self._create_oauth_client()
+
+    def _get_identity(self, provider_type: str | None) -> Identity:
+        """Retrieve the appropriate identity from the authentication context."""
+        identities = [x for x in self.auth_ctx.identities if x.provider_identity_id is not None]
+
+        if not identities:
+            raise ValueError("No identities found in authorization context.")
+
+        if provider_type is None:
+            if len(identities) > 1:
+                raise ValueError(
+                    "Multiple identities found. Please specify 'provider_type' parameter."
+                )
+            return identities[0]
+
+        identity = next((id for id in identities if id.provider_type == provider_type), None)
+
+        if identity is None:
+            raise ValueError(f"No identity found for provider '{provider_type}'.")
+
+        return identity
+
+    async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
+        """Get OAuth access token using the specified method."""
+        if auth_type != ToolAuth.OBO:
+            raise ValueError(
+                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+            )
+
+        identity = self._get_identity(provider_type)
+        token_data = await self.oauth_client.refresh_access_token(
+            identity_id=identity.provider_identity_id
+        )
+        return token_data.access_token
+
+    def _create_oauth_client(self) -> AsyncOAuthComponent:
+        """Create either DataRobot or Authlib OAuth client based on
+        authorization context.
+
+        Note: at the moment, only DataRobot OAuth client is supported.
+        """
+        logger.debug("Using DataRobot OAuth client")
+        return DatarobotAsyncOAuthClient()

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

@@ -92,8 +92,6 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
 
         # Use MCP context manager to handle connection lifecycle
         with mcp_tools_context(
-            api_base=self.api_base,
-            api_key=self.api_key,
             authorization_context=self._authorization_context,
         ) as mcp_tools:
             # Set MCP tools for all agents if MCP is not configured this is effectively a no-op

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

@@ -29,15 +29,10 @@ from datarobot_genai.core.mcp.common import MCPConfig
 
 @contextmanager
 def mcp_tools_context(
-    api_base: str | None = None,
-    api_key: str | None = None,
     authorization_context: dict[str, Any] | None = None,
 ) -> Generator[list[Any], None, None]:
     """Context manager for MCP tools that handles connection lifecycle."""
-    config = MCPConfig(
-        api_base=api_base, api_key=api_key, authorization_context=authorization_context
-    )
-
+    config = MCPConfig(authorization_context=authorization_context)
     # 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)
@@ -47,10 +42,8 @@ def mcp_tools_context(
     print(f"Connecting to MCP server: {config.server_config['url']}", flush=True)
 
     # Use MCPServerAdapter as context manager with the server config
-    adapter_setting = config.server_config.copy()
-    adapter_setting["transport"] = "streamable-http"
     try:
-        with MCPServerAdapter(adapter_setting) as tools:
+        with MCPServerAdapter(config.server_config) as tools:
             print(
                 f"Successfully connected to MCP server, got {len(tools)} tools",
                 flush=True,