datarobot-genai 0.2.39__py3-none-any.whl → 0.3.1__py3-none-any.whl

datarobot_genai/core/agents/__init__.py

@@ -17,7 +17,7 @@ This package provides:
 - BaseAgent: common initialization for agent env/config fields
 - Common helpers: make_system_prompt, extract_user_prompt_content
 - Framework utilities (optional extras):
-  - crewai: build_llm, create_pipeline_interactions_from_messages
+  - crewai: create_pipeline_interactions_from_messages
   - langgraph: create_pipeline_interactions_from_events
   - llamaindex: DataRobotLiteLLM, create_pipeline_interactions_from_events
 """

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/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/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/crewai/__init__.py

@@ -2,22 +2,19 @@
 
 Public API:
 - mcp_tools_context: Context manager returning available MCP tools for CrewAI.
-- build_llm: Construct a CrewAI LLM configured for DataRobot endpoints.
 - create_pipeline_interactions_from_messages: Convert messages to MultiTurnSample.
 """
 
 from datarobot_genai.core.mcp.common import MCPConfig
 
-from .agent import build_llm
+from .agent import CrewAIAgent
 from .agent import create_pipeline_interactions_from_messages
-from .base import CrewAIAgent
 from .events import CrewAIEventListener
 from .mcp import mcp_tools_context
 
 __all__ = [
     "mcp_tools_context",
     "CrewAIAgent",
-    "build_llm",
     "create_pipeline_interactions_from_messages",
     "CrewAIEventListener",
     "MCPConfig",

datarobot_genai/crewai/agent.py

@@ -12,26 +12,42 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from crewai import LLM
-from ragas import MultiTurnSample
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolMessage
+"""
+Base class for CrewAI-based agents.
 
-from datarobot_genai.core.utils.urls import get_api_base
+Manages MCP tool lifecycle and standardizes kickoff flow.
 
+Note: This base does not capture pipeline interactions; it returns None by
+default. Subclasses may implement message capture if they need interactions.
+"""
 
-def build_llm(
-    *,
-    api_base: str,
-    api_key: str | None,
-    model: str,
-    deployment_id: str | None,
-    timeout: int,
-) -> LLM:
-    """Create a CrewAI LLM configured for DataRobot LLM Gateway or deployment."""
-    base = get_api_base(api_base, deployment_id)
-    return LLM(model=model, api_base=base, api_key=api_key, timeout=timeout)
+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 datarobot_genai.core.agents.base import BaseAgent
+from datarobot_genai.core.agents.base import InvokeReturn
+from datarobot_genai.core.agents.base import UsageMetrics
+from datarobot_genai.core.agents.base import default_usage_metrics
+from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
+
+from .mcp import mcp_tools_context
+
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
 
 
 def create_pipeline_interactions_from_messages(
@@ -39,4 +55,121 @@ 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)
+
+
+class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
+    """Abstract base agent for CrewAI workflows.
+
+    Subclasses should define the ``agents`` and ``tasks`` properties
+    and may override ``build_crewai_workflow`` to customize the workflow
+    construction.
+    """
+
+    @property
+    @abc.abstractmethod
+    def agents(self) -> list[Any]:  # CrewAI Agent list
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def tasks(self) -> list[Any]:  # CrewAI Task list
+        raise NotImplementedError
+
+    def build_crewai_workflow(self) -> Any:
+        """Create a CrewAI workflow instance.
+
+        Default implementation constructs a Crew with provided agents and tasks.
+        Subclasses can override to customize Crew options.
+        """
+        return Crew(agents=self.agents, tasks=self.tasks, verbose=self.verbose)
+
+    @abc.abstractmethod
+    def make_kickoff_inputs(self, user_prompt_content: str) -> dict[str, Any]:
+        """Build the inputs dict for ``Crew.kickoff``.
+
+        Subclasses must implement this to provide the exact inputs required
+        by their CrewAI tasks.
+        """
+        raise NotImplementedError
+
+    def _extract_pipeline_interactions(self) -> MultiTurnSample | None:
+        """Extract pipeline interactions from event listener if available."""
+        if not hasattr(self, "event_listener"):
+            return None
+        try:
+            listener = getattr(self, "event_listener", None)
+            messages = getattr(listener, "messages", None) if listener is not None else None
+            return create_pipeline_interactions_from_messages(messages)
+        except Exception:
+            return None
+
+    def _extract_usage_metrics(self, crew_output: Any) -> UsageMetrics:
+        """Extract usage metrics from crew output."""
+        token_usage = getattr(crew_output, "token_usage", None)
+        if token_usage is not None:
+            return {
+                "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
+                "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
+                "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
+            }
+        return default_usage_metrics()
+
+    def _process_crew_output(
+        self, crew_output: Any
+    ) -> tuple[str, MultiTurnSample | None, UsageMetrics]:
+        """Process crew output into response tuple."""
+        response_text = str(crew_output.raw)
+        pipeline_interactions = self._extract_pipeline_interactions()
+        usage_metrics = self._extract_usage_metrics(crew_output)
+        return response_text, pipeline_interactions, usage_metrics
+
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
+        """Run the CrewAI workflow with the provided completion parameters."""
+        user_prompt_content = extract_user_prompt_content(completion_create_params)
+        # Preserve prior template startup print for CLI parity
+        try:
+            print("Running agent with user prompt:", user_prompt_content, flush=True)
+        except Exception:
+            # Printing is best-effort; proceed regardless
+            pass
+
+        # Use MCP context manager to handle connection lifecycle
+        with mcp_tools_context(
+            authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
+        ) as mcp_tools:
+            # Set MCP tools for all agents if MCP is not configured this is effectively a no-op
+            self.set_mcp_tools(mcp_tools)
+
+            # If an event listener is provided by the subclass/template, register it
+            if hasattr(self, "event_listener") and CrewAIEventsBus is not None:
+                try:
+                    listener = getattr(self, "event_listener")
+                    setup_fn = getattr(listener, "setup_listeners", None)
+                    if callable(setup_fn):
+                        setup_fn(CrewAIEventsBus)
+                except Exception:
+                    # Listener is optional best-effort; proceed without failing invoke
+                    pass
+
+            crew = self.build_crewai_workflow()
+
+            if is_streaming(completion_create_params):
+
+                async def _gen() -> AsyncGenerator[
+                    tuple[str, MultiTurnSample | None, UsageMetrics]
+                ]:
+                    crew_output = await asyncio.to_thread(
+                        crew.kickoff,
+                        inputs=self.make_kickoff_inputs(user_prompt_content),
+                    )
+                    yield self._process_crew_output(crew_output)
+
+                return _gen()
+
+            crew_output = crew.kickoff(inputs=self.make_kickoff_inputs(user_prompt_content))
+            return self._process_crew_output(crew_output)

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/drmcp/__init__.py

@@ -19,7 +19,8 @@ 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.core.utils import filter_tools_by_tags
+from datarobot_genai.drmcp.test_utils.clients.dr_gateway import DRLLMGatewayMCPClient
 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
@@ -66,11 +67,12 @@ __all__ = [
     "extract_datarobot_runtime_param_payload",
     "extract_datarobot_dict_runtime_param_payload",
     "extract_datarobot_credential_runtime_param_payload",
+    "filter_tools_by_tags",
     # Test utilities
     "get_dr_mcp_server_url",
     "get_headers",
     "ete_test_mcp_session",
-    "OpenAILLMMCPClient",
+    "DRLLMGatewayMCPClient",
     "ETETestExpectations",
     "ToolBaseE2E",
     "ToolCallTestExpectations",

datarobot_genai/drmcp/core/config.py

@@ -76,7 +76,25 @@ class MCPToolConfig(BaseSettings):
             RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
             "ENABLE_MICROSOFT_GRAPH_TOOLS",
         ),
-        description="Enable/disable Sharepoint tools",
+        description="Enable/disable Microsoft Graph (Sharepoint/OneDrive) tools",
+    )
+
+    enable_perplexity_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PERPLEXITY_TOOLS",
+            "ENABLE_PERPLEXITY_TOOLS",
+        ),
+        description="Enable/disable Perplexity tools",
+    )
+
+    enable_tavily_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_TAVILY_TOOLS",
+            "ENABLE_TAVILY_TOOLS",
+        ),
+        description="Enable/disable Tavily search tools",
     )
 
     is_atlassian_oauth_provider_configured: bool = Field(
@@ -131,6 +149,8 @@ class MCPToolConfig(BaseSettings):
         "enable_confluence_tools",
         "enable_gdrive_tools",
         "enable_microsoft_graph_tools",
+        "enable_perplexity_tools",
+        "enable_tavily_tools",
         "is_atlassian_oauth_provider_configured",
         "is_google_oauth_provider_configured",
         "is_microsoft_oauth_provider_configured",