datarobot-genai 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

datarobot_genai/core/utils/auth.py

@@ -12,23 +12,17 @@
 # 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 DatarobotOAuthClient
-from datarobot.auth.exceptions import OAuthProviderNotFound
-from datarobot.auth.exceptions import OAuthValidationErr
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
 from datarobot.auth.identity import Identity
-from datarobot.auth.oauth import OAuthToken
+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 mypyc.ir.ops import Sequence
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
@@ -189,203 +183,52 @@ 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(provider_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:
-    """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.
-    """
+    """Manages OAuth access tokens using generic OAuth client."""
 
     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
-        config = OAuthConfig.from_auth_ctx(auth_ctx)
-        self._retriever = create_token_retriever(config)
+        self.oauth_client = self._create_oauth_client()
 
     def _get_identity(self, provider_type: str | None) -> Identity:
-        """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.")
+        """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(oauth_identities) > 1:
-                raise OAuthValidationErr(
-                    "Multiple OAuth providers found. Specify 'provider_type' parameter."
+            if len(identities) > 1:
+                raise ValueError(
+                    "Multiple identities found. Please specify 'provider_type' parameter."
                 )
-            return oauth_identities[0]
+            return identities[0]
+
+        identity = next((id for id in identities if id.provider_type == provider_type), None)
 
-        identity = next((i for i in oauth_identities if i.provider_type == provider_type), None)
         if identity is None:
-            raise OAuthValidationErr(f"No identity found for provider '{provider_type}'.")
+            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 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.
-        """
+        """Get OAuth access token using the specified method."""
         if auth_type != ToolAuth.OBO:
             raise ValueError(
-                f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
+                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
             )
 
         identity = self._get_identity(provider_type)
-        token_data = await self._retriever.refresh_access_token(identity)
+        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/drmcp/tools/clients/confluence.py

@@ -12,3 +12,185 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Async client for interacting with Confluence Cloud REST API.
+
+At the moment of creating this client, official Confluence SDK is not supporting async.
+"""
+
+import logging
+from http import HTTPStatus
+from typing import Any
+
+import httpx
+from pydantic import BaseModel
+from pydantic import Field
+
+from .atlassian import ATLASSIAN_API_BASE
+from .atlassian import get_atlassian_cloud_id
+
+logger = logging.getLogger(__name__)
+
+
+class ConfluencePage(BaseModel):
+    """Pydantic model for Confluence page."""
+
+    page_id: str = Field(..., description="The unique page ID")
+    title: str = Field(..., description="Page title")
+    space_id: str = Field(..., description="Space ID where the page resides")
+    space_key: str | None = Field(None, description="Space key (if available)")
+    body: str = Field(..., description="Page content in storage format (HTML-like)")
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the page."""
+        return {
+            "page_id": self.page_id,
+            "title": self.title,
+            "space_id": self.space_id,
+            "space_key": self.space_key,
+            "body": self.body,
+        }
+
+
+class ConfluenceClient:
+    """
+    Client for interacting with Confluence API using OAuth access token.
+
+    At the moment of creating this client, official Confluence SDK is not supporting async.
+    """
+
+    EXPAND_FIELDS = "body.storage,space"
+
+    def __init__(self, access_token: str) -> None:
+        """
+        Initialize Confluence client with access token.
+
+        Args:
+            access_token: OAuth access token for Atlassian API
+        """
+        self.access_token = access_token
+        self._client = httpx.AsyncClient(
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+            },
+            timeout=30.0,
+        )
+        self._cloud_id: str | None = None
+
+    async def _get_cloud_id(self) -> str:
+        """
+        Get the cloud ID for the authenticated Atlassian Confluence instance.
+
+        According to Atlassian OAuth 2.0 documentation, API calls should use:
+        https://api.atlassian.com/ex/confluence/{cloudId}/wiki/rest/api/...
+
+        Returns
+        -------
+            Cloud ID string
+
+        Raises
+        ------
+            ValueError: If cloud ID cannot be retrieved
+        """
+        if self._cloud_id:
+            return self._cloud_id
+
+        self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="confluence")
+        return self._cloud_id
+
+    def _parse_response(self, data: dict) -> ConfluencePage:
+        """Parse API response into ConfluencePage."""
+        body_content = ""
+        body = data.get("body", {})
+        if isinstance(body, dict):
+            storage = body.get("storage", {})
+            if isinstance(storage, dict):
+                body_content = storage.get("value", "")
+
+        space = data.get("space", {})
+        space_key = space.get("key") if isinstance(space, dict) else None
+        space_id = space.get("id", "") if isinstance(space, dict) else data.get("spaceId", "")
+
+        return ConfluencePage(
+            page_id=str(data.get("id", "")),
+            title=data.get("title", ""),
+            space_id=str(space_id),
+            space_key=space_key,
+            body=body_content,
+        )
+
+    async def get_page_by_id(self, page_id: str) -> ConfluencePage:
+        """
+        Get a Confluence page by its ID.
+
+        Args:
+            page_id: The numeric page ID
+
+        Returns
+        -------
+            ConfluencePage with page data
+
+        Raises
+        ------
+            ValueError: If page is not found
+            httpx.HTTPStatusError: If the API request fails
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content/{page_id}"
+
+        response = await self._client.get(url, params={"expand": self.EXPAND_FIELDS})
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            raise ValueError(f"Page with ID '{page_id}' not found")
+
+        response.raise_for_status()
+        return self._parse_response(response.json())
+
+    async def get_page_by_title(self, title: str, space_key: str) -> ConfluencePage:
+        """
+        Get a Confluence page by its title within a specific space.
+
+        Args:
+            title: The exact page title
+            space_key: The space key where the page resides
+
+        Returns
+        -------
+            ConfluencePage with page data
+
+        Raises
+        ------
+            ValueError: If the page is not found
+            httpx.HTTPStatusError: If the API request fails
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content"
+
+        response = await self._client.get(
+            url,
+            params={
+                "title": title,
+                "spaceKey": space_key,
+                "expand": self.EXPAND_FIELDS,
+            },
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        results = data.get("results", [])
+
+        if not results:
+            raise ValueError(f"Page with title '{title}' not found in space '{space_key}'")
+
+        return self._parse_response(results[0])
+
+    async def __aenter__(self) -> "ConfluenceClient":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        """Async context manager exit."""
+        await self._client.aclose()

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -11,3 +11,71 @@
 # 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.
+
+"""Confluence MCP tools for interacting with Confluence Cloud."""
+
+import logging
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.atlassian import get_atlassian_access_token
+from datarobot_genai.drmcp.tools.clients.confluence import ConfluenceClient
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"confluence", "read", "get", "page"})
+async def confluence_get_page(
+    *,
+    page_id_or_title: Annotated[str, "The ID or the exact title of the Confluence page."],
+    space_key: Annotated[
+        str | None,
+        "Required if identifying the page by title. The space key (e.g., 'PROJ').",
+    ] = None,
+) -> ToolResult | ToolError:
+    """Retrieve the content of a specific Confluence page.
+
+    Use this tool to fetch Confluence pages by their numeric ID or by title.
+    Returns page content in HTML storage format.
+
+    Usage:
+        - By ID: page_id_or_title="856391684"
+        - By title: page_id_or_title="Meeting Notes", space_key="TEAM"
+
+    When using a page title, the space_key parameter is required.
+    """
+    if not page_id_or_title:
+        raise ToolError("Argument validation error: 'page_id_or_title' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with ConfluenceClient(access_token) as client:
+            if page_id_or_title.isdigit():
+                page_response = await client.get_page_by_id(page_id_or_title)
+            else:
+                if not space_key:
+                    raise ToolError(
+                        "Argument validation error: "
+                        "'space_key' is required when identifying a page by title."
+                    )
+                page_response = await client.get_page_by_title(page_id_or_title, space_key)
+    except ValueError as e:
+        logger.error(f"Value error getting Confluence page: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error getting Confluence page: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while getting Confluence page "
+            f"'{page_id_or_title}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Successfully retrieved page '{page_response.title}'.",
+        structured_content=page_response.as_flat_dict(),
+    )

{datarobot_genai-0.2.8.dist-info → datarobot_genai-0.2.10.dist-info}/METADATA

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

{datarobot_genai-0.2.8.dist-info → datarobot_genai-0.2.10.dist-info}/RECORD

@@ -15,7 +15,7 @@ datarobot_genai/core/cli/agent_kernel.py,sha256=3XX58DQ6XPpWB_tn5m3iGb3XTfhZf5X3
 datarobot_genai/core/mcp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/core/mcp/common.py,sha256=Y8SjuquUODKEfI7T9X-QuTMKdIlpCWFI1b3xs6tmHFA,7812
 datarobot_genai/core/utils/__init__.py,sha256=VxtRUz6iwb04eFQQy0zqTNXLAkYpPXcJxVoKV0nOdXk,59
-datarobot_genai/core/utils/auth.py,sha256=vvPYGmqJBkbx8FT-iOLibA9WafpIwhtwIRV92e3tLYI,14287
+datarobot_genai/core/utils/auth.py,sha256=LpSoHdPD2siskYwG8q4f9cike4VQdgFWJpuJrpiszXU,8674
 datarobot_genai/core/utils/urls.py,sha256=tk0t13duDEPcmwz2OnS4vwEdatruiuX8lnxMMhSaJik,2289
 datarobot_genai/crewai/__init__.py,sha256=MtFnHA3EtmgiK_GjwUGPgQQ6G1MCEzz1SDBwQi9lE8M,706
 datarobot_genai/crewai/agent.py,sha256=vp8_2LExpeLls7Fpzo0R6ud5I6Ryfu3n3oVTN4Yyi6A,1417
@@ -76,11 +76,11 @@ datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdY
 datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
-datarobot_genai/drmcp/tools/clients/confluence.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
+datarobot_genai/drmcp/tools/clients/confluence.py,sha256=gbVxeBe7RDEEQt5UMGGW6GoAXsYLhL009dOejYIaIiQ,6325
 datarobot_genai/drmcp/tools/clients/jira.py,sha256=JjvssdMAWgZ3HWZkQg0a3HjpE7yz7jfRtzO4LOp47Uw,3080
 datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
 datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
-datarobot_genai/drmcp/tools/confluence/tools.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
+datarobot_genai/drmcp/tools/confluence/tools.py,sha256=t5OqXIhUm6y9bAWymyqwEMElwTxGw1xRnkW2MgJrNF8,3106
 datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/jira/tools.py,sha256=EBf8T_VwKXNticorRN8ut9ktEMS0XFeb0Uj_6TKwMio,2191
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
@@ -105,9 +105,9 @@ datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZ
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
 datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
-datarobot_genai-0.2.8.dist-info/METADATA,sha256=vhjkIDaML9GN6yQ7X4On2edWKGO0KPWkvqDSvYpUShw,6172
-datarobot_genai-0.2.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.8.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.8.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.8.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.8.dist-info/RECORD,,
+datarobot_genai-0.2.10.dist-info/METADATA,sha256=AYl5uILzc_myorpYgMkAhbm6BZIGQaaHuMNQik6rdCE,6173
+datarobot_genai-0.2.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.10.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.10.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.10.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.10.dist-info/RECORD,,