datarobot-genai 0.2.39__tar.gz → 0.2.42__tar.gz

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/PKG-INFO

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

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.39"
+version = "0.2.42"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -130,6 +130,7 @@ dev = [
   "pytest-asyncio>=0.23.5",
   "responses>=0.25.8",
   "ipdb>=0.13.13",
+  "aioresponses>=0.7.8",
 ]
 
 
@@ -186,6 +187,9 @@ fixable = [
     "I001",   # isort import order within functions
     "PLC0415",# pylint import-outside-toplevel for controlled local imports
 ]
+"src/datarobot_genai/**/*.py" = [
+    "PLC0415",# pylint import-outside-toplevel: allow lazy imports for memory optimization
+]
 
 [tool.ruff.lint.isort]
 force-single-line = true

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/core/agents/base.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import abc
 import json
 import os
@@ -19,6 +21,7 @@ from collections.abc import AsyncGenerator
 from collections.abc import Mapping
 from typing import Any
 from typing import Generic
+from typing import TypeAlias
 from typing import TypedDict
 from typing import TypeVar
 from typing import cast
@@ -98,7 +101,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
         return get_api_base(self.api_base, deployment_id)
 
     @abc.abstractmethod
-    async def invoke(self, completion_create_params: CompletionCreateParams) -> "InvokeReturn":
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
         raise NotImplementedError("Not implemented")
 
     @classmethod
@@ -167,7 +170,7 @@ class UsageMetrics(TypedDict):
 
 
 # Canonical return type for DRUM-compatible invoke implementations
-InvokeReturn = (
+InvokeReturn: TypeAlias = (
     AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
     | tuple[str, MultiTurnSample | None, UsageMetrics]
 )

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/core/chat/responses.py

@@ -14,6 +14,8 @@
 
 """OpenAI-compatible response helpers for chat interactions."""
 
+from __future__ import annotations
+
 import asyncio
 import queue
 import time
@@ -24,6 +26,7 @@ from collections.abc import AsyncGenerator
 from collections.abc import AsyncIterator
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import TypeVar
 
@@ -38,10 +41,12 @@ from openai.types.chat import ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from openai.types.chat.chat_completion_chunk import ChoiceDelta
-from ragas import MultiTurnSample
 
 from datarobot_genai.core.agents import default_usage_metrics
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+
 
 class CustomModelChatResponse(ChatCompletion):
     pipeline_interactions: str | None = None

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/core/utils/auth.py

@@ -12,13 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import os
 import warnings
+from collections.abc import Sequence
 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
@@ -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.39 → datarobot_genai-0.2.42}/src/datarobot_genai/crewai/agent.py

@@ -11,15 +11,20 @@
 # 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from crewai import LLM
-from ragas import MultiTurnSample
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolMessage
 
 from datarobot_genai.core.utils.urls import get_api_base
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
+
 
 def build_llm(
     *,
@@ -39,4 +44,7 @@ def create_pipeline_interactions_from_messages(
 ) -> MultiTurnSample | None:
     if not messages:
         return None
+    # Lazy import to reduce memory overhead when ragas is not used
+    from ragas import MultiTurnSample
+
     return MultiTurnSample(user_input=messages)

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/crewai/base.py

@@ -26,13 +26,13 @@ from __future__ import annotations
 import abc
 import asyncio
 from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 from typing import Any
 
 from crewai import Crew
 from crewai.events.event_bus import CrewAIEventsBus
 from crewai.tools import BaseTool
 from openai.types.chat import CompletionCreateParams
-from ragas import MultiTurnSample
 
 from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
@@ -44,6 +44,9 @@ from datarobot_genai.core.agents.base import is_streaming
 from .agent import create_pipeline_interactions_from_messages
 from .mcp import mcp_tools_context
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+
 
 class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
     """Abstract base agent for CrewAI workflows.

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/crewai/events.py

@@ -16,12 +16,13 @@ from __future__ import annotations
 import importlib
 import json
 import logging
+from typing import TYPE_CHECKING
 from typing import Any
 
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolCall
-from ragas.messages import ToolMessage
+if TYPE_CHECKING:
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
 
 # Resolve crewai symbols at runtime to avoid mypy issues with untyped packages
 try:
@@ -65,6 +66,12 @@ class CrewAIEventListener:
         self.messages: list[HumanMessage | AIMessage | ToolMessage] = []
 
     def setup_listeners(self, crewai_event_bus: Any) -> None:
+        # Lazy import to reduce memory overhead when ragas is not used
+        from ragas.messages import AIMessage
+        from ragas.messages import HumanMessage
+        from ragas.messages import ToolCall
+        from ragas.messages import ToolMessage
+
         @crewai_event_bus.on(CrewKickoffStartedEvent)
         def on_crew_execution_started(_: Any, event: Any) -> None:
             self.messages.append(

{datarobot_genai-0.2.39 → datarobot_genai-0.2.42}/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -586,6 +586,102 @@ class MicrosoftGraphClient:
                 f"Microsoft Graph API error {response.status_code}: {response.text}"
             )
 
+    async def update_item_metadata(
+        self,
+        item_id: str,
+        fields_to_update: dict[str, Any],
+        site_id: str | None = None,
+        list_id: str | None = None,
+        drive_id: str | None = None,
+    ) -> dict[str, Any]:
+        """Update metadata on a SharePoint list item or OneDrive/SharePoint drive item.
+
+        For SharePoint list items: Updates custom column values via the fields endpoint.
+        For drive items: Updates properties like name and description.
+
+        Args:
+            item_id: The ID of the item to update.
+            fields_to_update: Key-value pairs of metadata fields to modify.
+            site_id: For SharePoint list items - the site ID.
+            list_id: For SharePoint list items - the list ID.
+            drive_id: For OneDrive/drive items - the drive ID.
+
+        Returns
+        -------
+            The API response containing updated item data.
+
+        Raises
+        ------
+            MicrosoftGraphError: If validation fails or the API request fails.
+        """
+        if not item_id or not item_id.strip():
+            raise MicrosoftGraphError("item_id cannot be empty")
+        if not fields_to_update:
+            raise MicrosoftGraphError("fields_to_update cannot be empty")
+
+        # Determine the endpoint based on provided parameters
+        has_sharepoint_context = site_id is not None and list_id is not None
+        has_drive_context = drive_id is not None
+
+        if has_sharepoint_context and has_drive_context:
+            raise MicrosoftGraphError(
+                "Cannot specify both SharePoint (site_id + list_id) and OneDrive "
+                "(document_library_id) context. Choose one."
+            )
+
+        if not has_sharepoint_context and not has_drive_context:
+            raise MicrosoftGraphError(
+                "Must specify either SharePoint context (site_id + list_id) or "
+                "OneDrive context (document_library_id)."
+            )
+
+        try:
+            if has_sharepoint_context:
+                # PATCH /sites/{site-id}/lists/{list-id}/items/{item-id}/fields
+                url = f"{GRAPH_API_BASE}/sites/{site_id}/lists/{list_id}/items/{item_id}/fields"
+                response = await self._client.patch(url, json=fields_to_update)
+            else:
+                # Drive item: PATCH /drives/{drive-id}/items/{item-id}
+                url = f"{GRAPH_API_BASE}/drives/{drive_id}/items/{item_id}"
+                response = await self._client.patch(url, json=fields_to_update)
+
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_update_metadata_error(e, item_id) from e
+
+    def _handle_update_metadata_error(
+        self,
+        error: httpx.HTTPStatusError,
+        item_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for metadata updates and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to update metadata: HTTP {status_code}"
+
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                api_message = error_data.get("error", {}).get("message", "Invalid request")
+                error_msg = f"Bad request updating metadata: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid field names or values."
+        elif status_code == 401:
+            error_msg = "Authentication failed. Access token may be expired or invalid."
+        elif status_code == 403:
+            error_msg = (
+                f"Permission denied: you don't have permission to update item '{item_id}'. "
+                "Requires Sites.ReadWrite.All or Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = f"Item '{item_id}' not found."
+        elif status_code == 409:
+            error_msg = f"Conflict: item '{item_id}' may have been modified concurrently."
+        elif status_code == 429:
+            error_msg = "Rate limit exceeded. Please try again later."
+
+        return MicrosoftGraphError(error_msg)
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self