datarobot-genai 0.2.7__tar.gz → 0.2.40__tar.gz

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.7
+Version: 0.2.40
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -25,6 +25,8 @@ 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: anthropic<1.0.0,~=0.71.0; extra == 'crewai'
+Requires-Dist: azure-ai-inference<2.0.0,>=1.0.0b9; extra == 'crewai'
 Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
 Requires-Dist: crewai>=1.1.0; extra == 'crewai'
 Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.7"
+version = "0.2.40"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -40,6 +40,8 @@ datarobot_auth_provider = "datarobot_genai.nat.datarobot_auth_provider"
 
 [project.optional-dependencies]
 crewai = [
+  "anthropic~=0.71.0,<1.0.0",  # Needed for integration with anthropic endpoints
+  "azure-ai-inference>=1.0.0b9,<2.0.0",  # Needed for integration with azure endpoints
   "crewai>=1.1.0",
   "crewai-tools[mcp]>=0.69.0,<0.77.0",
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
@@ -128,6 +130,7 @@ dev = [
   "pytest-asyncio>=0.23.5",
   "responses>=0.25.8",
   "ipdb>=0.13.13",
+  "aioresponses>=0.7.8",
 ]
 
 
@@ -203,6 +206,9 @@ asyncio_mode = "auto"
 source = ["datarobot_genai"]
 omit = [
     "*/__init__.py",
+    "*/test_utils/*",
+    # nat requires Python >=3.11, can't be imported on 3.10 for coverage
+    "*/nat/*",
 ]
 
 [tool.coverage.report]

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/core/cli/agent_kernel.py

@@ -107,7 +107,10 @@ class AgentKernel:
         stream: bool = False,
         config: Any | None = None,
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
-        chat_api_url = config.agent_endpoint if config else self.base_url
+        if config is not None:
+            chat_api_url = f"http://localhost:{config.local_dev_port}"
+        else:
+            chat_api_url = self.base_url
         print(chat_api_url)
 
         return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/core/utils/auth.py

@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import os
 import warnings
 from typing import Any
+from typing import Protocol
 
+import aiohttp
 import jwt
-from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotOAuthClient
+from datarobot.auth.exceptions import OAuthProviderNotFound
+from datarobot.auth.exceptions import OAuthValidationErr
 from datarobot.auth.identity import Identity
-from datarobot.auth.oauth import AsyncOAuthComponent
+from datarobot.auth.oauth import OAuthToken
 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 mypyc.ir.ops import Sequence
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
@@ -183,52 +189,203 @@ class AuthContextHeaderHandler:
             return None
 
 
+# --- OAuth Token Provider Implementation ---
+
+
+class TokenRetriever(Protocol):
+    """Protocol for OAuth token retrievers."""
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only those valid for this retriever implementation."""
+        ...
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token for the given identity ID.
+
+        Parameters
+        ----------
+        identity_id : str
+            The provider identity ID to refresh the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The refreshed OAuth token.
+        """
+        ...
+
+
+class DatarobotTokenRetriever:
+    """Retrieves OAuth tokens using the DataRobot platform."""
+
+    def __init__(self) -> None:
+        self._client = DatarobotOAuthClient()
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter oauth2 identities to only those with provider_identity_id.
+
+        The `provider_identity_id` is required in order to identify the provider
+        and retrieve the access token from DataRobot OAuth Providers service.
+        """
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token using DataRobot's OAuth client."""
+        return await self._client.refresh_access_token(identity_id=identity.provider_identity_id)
+
+
+class AuthlibTokenRetriever:
+    """Retrieves OAuth tokens from a generic Authlib-based endpoint."""
+
+    def __init__(self, application_endpoint: str) -> None:
+        if not application_endpoint:
+            raise ValueError("AuthlibTokenRetriever requires 'application_endpoint'.")
+        self.application_endpoint = application_endpoint.rstrip("/")
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only OAuth2 identities."""
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id is None]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Retrieve an OAuth token via an HTTP POST request.
+
+        Parameters
+        ----------
+        identity : Identity
+            The identity to retrieve the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The retrieved OAuth token.
+        """
+        api_token = os.environ.get("DATAROBOT_API_TOKEN")
+        if not api_token:
+            raise ValueError("DATAROBOT_API_TOKEN environment variable is required but not set.")
+
+        token_url = f"{self.application_endpoint}/oauth/token/"
+        headers = {"Authorization": f"Bearer {api_token}"}
+        payload = {"identity_id": identity.id}
+        timeout = aiohttp.ClientTimeout(total=30)
+
+        try:
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(token_url, headers=headers, json=payload) as response:
+                    response.raise_for_status()
+                    data = await response.json()
+                    logger.debug(f"Retrieved access token from {token_url}")
+                    return OAuthToken(**data)
+        except aiohttp.ClientError as e:
+            logger.error(f"Error retrieving token from {token_url}: {e}")
+            raise
+
+
+class OAuthConfig(BaseModel):
+    """Configuration extracted from AuthCtx metadata for OAuth operations."""
+
+    implementation: str = "datarobot"
+    application_endpoint: str | None = None
+
+    @classmethod
+    def from_auth_ctx(cls, auth_ctx: AuthCtx) -> "OAuthConfig":
+        metadata = auth_ctx.metadata or {}
+        return cls(
+            implementation=metadata.get("oauth_implementation", "datarobot"),
+            application_endpoint=metadata.get("application_endpoint"),
+        )
+
+
+def create_token_retriever(config: OAuthConfig) -> TokenRetriever:
+    """Create a token retriever based on the OAuth configuration.
+
+    Parameters
+    ----------
+    config : OAuthConfig
+        The OAuth configuration specifying implementation type and endpoints.
+
+    Returns
+    -------
+    TokenRetriever
+        The configured token retriever instance.
+    """
+    if config.implementation == "datarobot":
+        return DatarobotTokenRetriever()
+
+    if config.implementation == "authlib":
+        if not config.application_endpoint:
+            raise ValueError("Required 'application_endpoint' not found in metadata.")
+        return AuthlibTokenRetriever(config.application_endpoint)
+
+    raise ValueError(
+        f"Unsupported OAuth implementation: '{config.implementation}'. "
+        f"Supported values: datarobot, authlib."
+    )
+
+
 class AsyncOAuthTokenProvider:
-    """Manages OAuth access tokens using generic OAuth client."""
+    """Provides OAuth tokens for authorized users.
+
+    This class manages OAuth token retrieval for users with multiple identity providers.
+    It uses either DataRobot or Authlib as the OAuth token storage and refresh backend
+    based on the auth context metadata.
+    """
 
     def __init__(self, auth_ctx: AuthCtx) -> None:
+        """Initialize the provider with an authorization context.
+
+        Parameters
+        ----------
+        auth_ctx : AuthCtx
+            The authorization context containing user identities and metadata.
+        """
         self.auth_ctx = auth_ctx
-        self.oauth_client = self._create_oauth_client()
+        config = OAuthConfig.from_auth_ctx(auth_ctx)
+        self._retriever = create_token_retriever(config)
 
     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.")
+        """Get identity from auth context, filtered by provider_type if specified."""
+        oauth_identities = self._retriever.filter_identities(self.auth_ctx.identities)
+        if not oauth_identities:
+            raise OAuthProviderNotFound("No OAuth provider found.")
 
         if provider_type is None:
-            if len(identities) > 1:
-                raise ValueError(
-                    "Multiple identities found. Please specify 'provider_type' parameter."
+            if len(oauth_identities) > 1:
+                raise OAuthValidationErr(
+                    "Multiple OAuth providers found. Specify 'provider_type' parameter."
                 )
-            return identities[0]
-
-        identity = next((id for id in identities if id.provider_type == provider_type), None)
+            return oauth_identities[0]
 
+        identity = next((i for i in oauth_identities if i.provider_type == provider_type), None)
         if identity is None:
-            raise ValueError(f"No identity found for provider '{provider_type}'.")
-
+            raise OAuthValidationErr(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."""
+        """Get an OAuth access token for the specified auth type and provider.
+
+        Parameters
+        ----------
+        auth_type : ToolAuth
+            Authentication type (only OBO is supported).
+        provider_type : str, optional
+            The specific provider to use (e.g., 'google'). Required if multiple
+            identities are available.
+
+        Returns
+        -------
+        str
+            The retrieved OAuth access token.
+
+        Raises
+        ------
+        ValueError
+            If the auth type is unsupported or if a suitable identity cannot be found.
+        """
         if auth_type != ToolAuth.OBO:
             raise ValueError(
-                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+                f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
             )
 
         identity = self._get_identity(provider_type)
-        token_data = await self.oauth_client.refresh_access_token(
-            identity_id=identity.provider_identity_id
-        )
+        token_data = await self._retriever.refresh_access_token(identity)
         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.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/drmcp/__init__.py

@@ -19,11 +19,11 @@ A reusable library for building Model Context Protocol (MCP) servers with DataRo
 """
 
 # Export main server components
+from datarobot_genai.drmcp.test_utils.clients.openai import OpenAILLMMCPClient
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import ete_test_mcp_session
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_dr_mcp_server_url
 from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_headers
 from datarobot_genai.drmcp.test_utils.mcp_utils_integration import integration_test_mcp_session
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ETETestExpectations
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolBaseE2E
 from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolCallTestExpectations
@@ -70,7 +70,7 @@ __all__ = [
     "get_dr_mcp_server_url",
     "get_headers",
     "ete_test_mcp_session",
-    "LLMMCPClient",
+    "OpenAILLMMCPClient",
     "ETETestExpectations",
     "ToolBaseE2E",
     "ToolCallTestExpectations",

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/drmcp/core/config.py

@@ -31,6 +31,124 @@ from .constants import DEFAULT_DATAROBOT_ENDPOINT
 from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
 
 
+class MCPToolConfig(BaseSettings):
+    """Tool configuration for MCP server."""
+
+    enable_predictive_tools: bool = Field(
+        default=True,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
+            "ENABLE_PREDICTIVE_TOOLS",
+        ),
+        description="Enable/disable predictive tools",
+    )
+
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+
+    enable_microsoft_graph_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
+            "ENABLE_MICROSOFT_GRAPH_TOOLS",
+        ),
+        description="Enable/disable Sharepoint tools",
+    )
+
+    is_atlassian_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+            "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Atlassian OAuth provider is configured for Atlassian integration",
+    )
+
+    @property
+    def is_atlassian_oauth_configured(self) -> bool:
+        """Check if Atlassian OAuth is configured via provider flag or environment variables."""
+        return self.is_atlassian_oauth_provider_configured or bool(
+            os.getenv("ATLASSIAN_CLIENT_ID") and os.getenv("ATLASSIAN_CLIENT_SECRET")
+        )
+
+    is_google_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Google OAuth provider is configured for Google integration",
+    )
+
+    @property
+    def is_google_oauth_configured(self) -> bool:
+        return self.is_google_oauth_provider_configured or bool(
+            os.getenv("GOOGLE_CLIENT_ID") and os.getenv("GOOGLE_CLIENT_SECRET")
+        )
+
+    is_microsoft_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+            "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Microsoft OAuth provider is configured for Microsoft integration",
+    )
+
+    @property
+    def is_microsoft_oauth_configured(self) -> bool:
+        return self.is_microsoft_oauth_provider_configured or bool(
+            os.getenv("MICROSOFT_CLIENT_ID") and os.getenv("MICROSOFT_CLIENT_SECRET")
+        )
+
+    @field_validator(
+        "enable_predictive_tools",
+        "enable_jira_tools",
+        "enable_confluence_tools",
+        "enable_gdrive_tools",
+        "enable_microsoft_graph_tools",
+        "is_atlassian_oauth_provider_configured",
+        "is_google_oauth_provider_configured",
+        "is_microsoft_oauth_provider_configured",
+        mode="before",
+    )
+    @classmethod
+    def validate_runtime_params(cls, v: Any) -> Any:
+        """Validate runtime parameters."""
+        return extract_datarobot_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
 class MCPServerConfig(BaseSettings):
     """MCP Server configuration using pydantic settings."""
 
@@ -188,62 +306,11 @@ class MCPServerConfig(BaseSettings):
         ),
         description="Enable/disable memory management",
     )
-    enable_predictive_tools: bool = Field(
-        default=True,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
-            "ENABLE_PREDICTIVE_TOOLS",
-        ),
-        description="Enable/disable predictive tools",
-    )
 
-    # Jira tools
-    enable_jira_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
-            "ENABLE_JIRA_TOOLS",
-        ),
-        description="Enable/disable Jira tools",
-    )
-    is_jira_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Jira OAuth provider is configured for Jira integration",
-    )
-
-    @property
-    def is_jira_oauth_configured(self) -> bool:
-        return self.is_jira_oauth_provider_configured or bool(
-            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
-        )
-
-    # Confluence tools
-    enable_confluence_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
-            "ENABLE_CONFLUENCE_TOOLS",
-        ),
-        description="Enable/disable Confluence tools",
+    tool_config: MCPToolConfig = Field(
+        default_factory=MCPToolConfig,
+        description="Tool configuration",
     )
-    is_confluence_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Confluence OAuth provider is configured for Confluence integration",
-    )
-
-    @property
-    def is_confluence_oauth_configured(self) -> bool:
-        return self.is_confluence_oauth_provider_configured or bool(
-            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
-        )
 
     @field_validator(
         "otel_attributes",
@@ -267,11 +334,6 @@ class MCPServerConfig(BaseSettings):
         "mcp_server_register_dynamic_tools_on_startup",
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
-        "enable_predictive_tools",
-        "enable_jira_tools",
-        "is_jira_oauth_provider_configured",
-        "enable_confluence_tools",
-        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -31,9 +31,6 @@ from .dynamic_prompts.register import register_prompts_from_datarobot_prompt_man
 from .dynamic_tools.deployment.register import register_tools_of_datarobot_deployments
 from .logging import MCPLogging
 from .mcp_instance import mcp
-from .mcp_server_tools import get_all_available_tags  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import get_tool_info_by_name  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import list_tools_by_tags  # noqa # pylint: disable=unused-import
 from .memory_management.manager import MemoryManager
 from .routes import register_routes
 from .routes_utils import prefix_mount_path

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/drmcp/core/exceptions.py

@@ -19,7 +19,3 @@ class DynamicToolRegistrationError(Exception):
 
 class DynamicPromptRegistrationError(Exception):
     """Exception raised for errors in the dynamic prompt registration process."""
-
-
-class MCPError(Exception):
-    """Base class for MCP errors."""

{datarobot_genai-0.2.7 → datarobot_genai-0.2.40}/src/datarobot_genai/drmcp/core/logging.py

@@ -20,7 +20,7 @@ from collections.abc import Callable
 from typing import Any
 from typing import TypeVar
 
-from .exceptions import MCPError
+from fastmcp.exceptions import ToolError
 
 # Secret patterns to redact from logs
 SECRET_PATTERNS = [
@@ -93,6 +93,6 @@ def log_execution(func: F) -> F:
             return result
         except Exception as e:
             error_msg = _log_error(logger, func.__name__, e, args=args, kwargs=kwargs)
-            raise MCPError(error_msg)
+            raise ToolError(error_msg)
 
     return wrapper  # type: ignore[return-value]