datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/core/utils/auth.py

@@ -0,0 +1,234 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import warnings
+from typing import Any
+
+import jwt
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.identity import Identity
+from datarobot.auth.oauth import AsyncOAuthComponent
+from datarobot.auth.session import AuthCtx
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from datarobot.models.genai.agent.auth import ToolAuth
+from datarobot.models.genai.agent.auth import get_authorization_context
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class AuthContextConfig(DataRobotAppFrameworkBaseSettings):
+    session_secret_key: str = ""
+
+
+class DRAppCtx(BaseModel):
+    """DataRobot application context from authorization metadata."""
+
+    email: str | None = None
+    api_key: str | None = None
+
+
+class AuthContextHeaderHandler:
+    """Manages encoding and decoding of authorization context into JWT tokens.
+
+    This class provides a consistent interface for encoding auth context into JWT tokens
+    and exchanging them via HTTP headers across multiple applications.
+    """
+
+    HEADER_NAME = "X-DataRobot-Authorization-Context"
+    DEFAULT_ALGORITHM = "HS256"
+
+    def __init__(
+        self,
+        secret_key: str | None = None,
+        algorithm: str = DEFAULT_ALGORITHM,
+        validate_signature: bool = True,
+    ) -> None:
+        """Initialize the handler.
+
+        Parameters
+        ----------
+        secret_key : Optional[str]
+            Secret key for JWT encoding/decoding. If None, tokens will be unsigned (insecure).
+        algorithm : str
+            JWT algorithm. Default is "HS256".
+        validate_signature : bool
+            Whether to validate JWT signatures. Default is True.
+
+        Raises
+        ------
+        ValueError
+            If algorithm is 'none' (insecure).
+        """
+        if algorithm is None:
+            raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")
+
+        # Get secret key from parameter, config, or environment variable
+        # Handle the case where AuthContextConfig() initialization fails due to
+        # a bug in the datarobot package when SESSION_SECRET_KEY is not set
+        if secret_key:
+            self.secret_key = secret_key
+        else:
+            try:
+                config = AuthContextConfig()
+                self.secret_key = config.session_secret_key or ""
+            except (TypeError, AttributeError, Exception):
+                # Fallback to reading environment variable directly if config initialization fails
+                # This can happen when SESSION_SECRET_KEY is not set and the datarobot package's
+                # getenv function encounters a bug with None values
+                # it tries to check if "apiToken" in payload: when payload is None
+                self.secret_key = ""
+
+        self.algorithm = algorithm
+        self.validate_signature = validate_signature
+
+    @property
+    def header(self) -> str:
+        """Get the header name for authorization context."""
+        return self.HEADER_NAME
+
+    def get_header(self, authorization_context: dict[str, Any] | None = None) -> dict[str, str]:
+        """Get the authorization context header with encoded JWT token."""
+        token = self.encode(authorization_context)
+        if not token:
+            return {}
+
+        return {self.header: token}
+
+    def encode(self, authorization_context: dict[str, Any] | None = None) -> str | None:
+        """Encode the current authorization context into a JWT token."""
+        auth_context = authorization_context or get_authorization_context()
+        if not auth_context:
+            return None
+
+        if not self.secret_key:
+            warnings.warn(
+                "No secret key provided. Please make sure SESSION_SECRET_KEY is set. "
+                "JWT tokens will be signed with an empty key. This is insecure and should "
+                "only be used for testing."
+            )
+
+        return jwt.encode(auth_context, self.secret_key, algorithm=self.algorithm)
+
+    def decode(self, token: str) -> dict[str, Any] | None:
+        """Decode a JWT token into the authorization context."""
+        if not token:
+            return None
+
+        if not self.secret_key and self.validate_signature:
+            logger.error(
+                "No secret key provided. Cannot validate signature. "
+                "Provide a secret key or set validate_signature to False."
+            )
+            return None
+
+        try:
+            decoded = jwt.decode(
+                jwt=token,
+                key=self.secret_key,
+                algorithms=[self.algorithm],
+                options={"verify_signature": self.validate_signature},
+            )
+        except jwt.ExpiredSignatureError:
+            logger.info("JWT token has expired.")
+            return None
+        except jwt.InvalidTokenError:
+            logger.warning("JWT token is invalid or malformed.")
+            return None
+
+        if not isinstance(decoded, dict):
+            logger.warning("Decoded JWT token is not a dictionary.")
+            return None
+
+        return decoded
+
+    def get_context(self, headers: dict[str, str]) -> AuthCtx | None:
+        """Extract and validate authorization context from headers.
+
+        Parameters
+        ----------
+        headers : Dict[str, str]
+            HTTP headers containing the authorization context.
+
+        Returns
+        -------
+        Optional[AuthCtx]
+            Validated authorization context or None if validation fails.
+        """
+        token = headers.get(self.header) or headers.get(self.header.lower())
+        if not token:
+            logger.debug("No authorization context header found")
+            return None
+
+        auth_ctx_dict = self.decode(token)
+        if not auth_ctx_dict:
+            logger.debug("Failed to decode auth context from token")
+            return None
+
+        try:
+            return AuthCtx(**auth_ctx_dict)
+        except Exception as e:
+            logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
+            return None
+
+
+class AsyncOAuthTokenProvider:
+    """Manages OAuth access tokens using generic OAuth client."""
+
+    def __init__(self, auth_ctx: AuthCtx) -> None:
+        self.auth_ctx = auth_ctx
+        self.oauth_client = self._create_oauth_client()
+
+    def _get_identity(self, provider_type: str | None) -> Identity:
+        """Retrieve the appropriate identity from the authentication context."""
+        identities = [x for x in self.auth_ctx.identities if x.provider_identity_id is not None]
+
+        if not identities:
+            raise ValueError("No identities found in authorization context.")
+
+        if provider_type is None:
+            if len(identities) > 1:
+                raise ValueError(
+                    "Multiple identities found. Please specify 'provider_type' parameter."
+                )
+            return identities[0]
+
+        identity = next((id for id in identities if id.provider_type == provider_type), None)
+
+        if identity is None:
+            raise ValueError(f"No identity found for provider '{provider_type}'.")
+
+        return identity
+
+    async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
+        """Get OAuth access token using the specified method."""
+        if auth_type != ToolAuth.OBO:
+            raise ValueError(
+                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+            )
+
+        identity = self._get_identity(provider_type)
+        token_data = await self.oauth_client.refresh_access_token(
+            identity_id=identity.provider_identity_id
+        )
+        return token_data.access_token
+
+    def _create_oauth_client(self) -> AsyncOAuthComponent:
+        """Create either DataRobot or Authlib OAuth client based on
+        authorization context.
+
+        Note: at the moment, only DataRobot OAuth client is supported.
+        """
+        logger.debug("Using DataRobot OAuth client")
+        return DatarobotAsyncOAuthClient()

datarobot_genai/core/utils/urls.py

@@ -0,0 +1,64 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from urllib.parse import urlparse
+from urllib.parse import urlunparse
+
+
+def get_api_base(api_base: str, deployment_id: str | None) -> str:
+    """
+    Construct the LiteLLM API base URL for a deployment.
+
+    Parameters
+    ----------
+    api_base : str
+        Base URL for the LiteLLM API.
+    deployment_id : str | None
+        Deployment identifier. When provided, a chat/completions URL is produced.
+
+    Returns
+    -------
+    str
+        Normalized URL for the given deployment. Ensures a trailing slash unless the path
+        ends with "chat/completions" or already has a meaningful path component.
+    """
+    # Normalize the URL and drop a trailing /api/v2 if present
+    parsed = urlparse(api_base)
+    path = re.sub(r"/api/v2/?$", "", parsed.path)
+    base_url = urlunparse(
+        (
+            parsed.scheme,
+            parsed.netloc,
+            path,
+            parsed.params,
+            parsed.query,
+            parsed.fragment,
+        )
+    )
+    base_url = base_url.rstrip("/")
+
+    # If the base_url already ends with chat/completions, return it.
+    if base_url.endswith("chat/completions"):
+        return base_url
+
+    # If the path contains deployments or genai, it's already a complete API path preserve it.
+    if path and ("deployments" in path or "genai" in path):
+        return f"{base_url}/" if not base_url.endswith("/") else base_url
+
+    # For all other cases (including custom base paths), apply deployment logic if needed.
+    if deployment_id:
+        return f"{base_url}/api/v2/deployments/{deployment_id}/chat/completions"
+    # Otherwise, just return the base URL with a trailing slash for normalization.
+    return f"{base_url}/"

datarobot_genai/crewai/__init__.py

@@ -0,0 +1,24 @@
+"""CrewAI utilities and helpers.
+
+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 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

@@ -0,0 +1,42 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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
+
+
+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)
+
+
+def create_pipeline_interactions_from_messages(
+    messages: list[HumanMessage | AIMessage | ToolMessage] | None,
+) -> MultiTurnSample | None:
+    if not messages:
+        return None
+    return MultiTurnSample(user_input=messages)

datarobot_genai/crewai/base.py

@@ -0,0 +1,159 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Base class for CrewAI-based agents.
+
+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.
+"""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+from collections.abc import AsyncGenerator
+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
+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 .agent import create_pipeline_interactions_from_messages
+from .mcp import mcp_tools_context
+
+
+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

@@ -0,0 +1,117 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+from typing import Any
+
+from ragas.messages import AIMessage
+from ragas.messages import HumanMessage
+from ragas.messages import ToolCall
+from ragas.messages import ToolMessage
+
+# Resolve crewai symbols at runtime to avoid mypy issues with untyped packages
+try:
+    _events_mod = importlib.import_module("crewai.events.event_types")
+    AgentExecutionCompletedEvent = getattr(_events_mod, "AgentExecutionCompletedEvent")
+    AgentExecutionStartedEvent = getattr(_events_mod, "AgentExecutionStartedEvent")
+    CrewKickoffStartedEvent = getattr(_events_mod, "CrewKickoffStartedEvent")
+    ToolUsageFinishedEvent = getattr(_events_mod, "ToolUsageFinishedEvent")
+    ToolUsageStartedEvent = getattr(_events_mod, "ToolUsageStartedEvent")
+
+    _bus_mod = importlib.import_module("crewai.events.event_bus")
+    CrewAIEventsBus = getattr(_bus_mod, "CrewAIEventsBus")
+
+    _base_mod = importlib.import_module("crewai.events.base_event_listener")
+    _RuntimeBaseEventListener = getattr(_base_mod, "BaseEventListener")
+except Exception:
+    try:  # pragma: no cover - compatibility for older crewai
+        _events_mod = importlib.import_module("crewai.utilities.events")
+        AgentExecutionCompletedEvent = getattr(_events_mod, "AgentExecutionCompletedEvent")
+        AgentExecutionStartedEvent = getattr(_events_mod, "AgentExecutionStartedEvent")
+        CrewKickoffStartedEvent = getattr(_events_mod, "CrewKickoffStartedEvent")
+        ToolUsageFinishedEvent = getattr(_events_mod, "ToolUsageFinishedEvent")
+        ToolUsageStartedEvent = getattr(_events_mod, "ToolUsageStartedEvent")
+
+        _bus_mod = importlib.import_module("crewai.utilities.events")
+        CrewAIEventsBus = getattr(_bus_mod, "CrewAIEventsBus")
+        _base_mod = importlib.import_module("crewai.utilities.events.base_event_listener")
+        _RuntimeBaseEventListener = getattr(_base_mod, "BaseEventListener")
+    except Exception:
+        raise ImportError(
+            "CrewAI is required for datarobot_genai.crewai.* modules. "
+            "Install with the CrewAI extra:\n"
+            "  install 'datarobot-genai[crewai]'"
+        )
+
+
+class CrewAIEventListener:
+    """Collects CrewAI events into Ragas messages for pipeline interactions."""
+
+    def __init__(self) -> None:
+        self.messages: list[HumanMessage | AIMessage | ToolMessage] = []
+
+    def setup_listeners(self, crewai_event_bus: Any) -> None:
+        @crewai_event_bus.on(CrewKickoffStartedEvent)
+        def on_crew_execution_started(_: Any, event: Any) -> None:
+            self.messages.append(
+                HumanMessage(content=f"Working on input '{json.dumps(event.inputs)}'")
+            )
+
+        @crewai_event_bus.on(AgentExecutionStartedEvent)
+        def on_agent_execution_started(_: Any, event: Any) -> None:
+            self.messages.append(AIMessage(content=event.task_prompt, tool_calls=[]))
+
+        @crewai_event_bus.on(AgentExecutionCompletedEvent)
+        def on_agent_execution_completed(_: Any, event: Any) -> None:
+            self.messages.append(AIMessage(content=event.output, tool_calls=[]))
+
+        @crewai_event_bus.on(ToolUsageStartedEvent)
+        def on_tool_usage_started(_: Any, event: Any) -> None:
+            # It's a tool call - add tool call to last AIMessage
+            if len(self.messages) == 0:
+                logging.warning("Direct tool usage without agent invocation")
+                return
+            last_message = self.messages[-1]
+            if not isinstance(last_message, AIMessage):
+                logging.warning(
+                    "Tool call must be preceded by an AIMessage somewhere in the conversation."
+                )
+                return
+            if isinstance(event.tool_args, (str, bytes, bytearray)):
+                parsed_args: Any = json.loads(event.tool_args)
+            else:
+                parsed_args = event.tool_args
+            tool_call = ToolCall(name=event.tool_name, args=parsed_args)
+            if last_message.tool_calls is None:
+                last_message.tool_calls = []
+            last_message.tool_calls.append(tool_call)
+
+        @crewai_event_bus.on(ToolUsageFinishedEvent)
+        def on_tool_usage_finished(_: Any, event: Any) -> None:
+            if len(self.messages) == 0:
+                logging.warning("Direct tool usage without agent invocation")
+                return
+            last_message = self.messages[-1]
+            if not isinstance(last_message, AIMessage):
+                logging.warning(
+                    "Tool call must be preceded by an AIMessage somewhere in the conversation."
+                )
+                return
+            if not last_message.tool_calls:
+                logging.warning("No previous tool calls found")
+                return
+            self.messages.append(ToolMessage(content=event.output))