datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/__init__.py

@@ -0,0 +1,19 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+
+from datarobot_genai.core.utils.urls import get_api_base
+
+"""
+Only add imports with the core dependencies here at the top level.
+For the optional extras, these need to be imported from their respective sub packages.
+"""
+
+__all__ = [
+    "get_api_base",
+    "__version__",
+]
+
+try:
+    __version__ = version("datarobot-genai")
+except PackageNotFoundError:  # pragma: no cover - during local dev without install
+    __version__ = "0.0.0"

datarobot_genai/core/agents/__init__.py

@@ -0,0 +1,43 @@
+# 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.
+"""Reusable agent utilities and base classes for end-user templates.
+
+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
+  - langgraph: create_pipeline_interactions_from_events
+  - llamaindex: DataRobotLiteLLM, create_pipeline_interactions_from_events
+"""
+
+from ..mcp.common import MCPConfig
+from .base import BaseAgent
+from .base import InvokeReturn
+from .base import UsageMetrics
+from .base import default_usage_metrics
+from .base import extract_user_prompt_content
+from .base import is_streaming
+from .base import make_system_prompt
+
+__all__ = [
+    "BaseAgent",
+    "make_system_prompt",
+    "extract_user_prompt_content",
+    "default_usage_metrics",
+    "is_streaming",
+    "InvokeReturn",
+    "UsageMetrics",
+    "MCPConfig",
+]

datarobot_genai/core/agents/base.py

@@ -0,0 +1,195 @@
+# 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 abc
+import json
+import os
+from collections.abc import AsyncGenerator
+from collections.abc import Mapping
+from typing import Any
+from typing import Generic
+from typing import TypedDict
+from typing import TypeVar
+from typing import cast
+
+from ag_ui.core import Event
+from openai.types.chat import CompletionCreateParams
+from ragas import MultiTurnSample
+
+from datarobot_genai.core.utils.urls import get_api_base
+
+TTool = TypeVar("TTool")
+
+
+class BaseAgent(Generic[TTool], abc.ABC):
+    """BaseAgent centralizes common initialization for agent templates.
+
+    Fields:
+      - api_key: DataRobot API token
+      - api_base: Endpoint for DataRobot, normalized for LLM Gateway usage
+      - model: Preferred model name
+      - timeout: Request timeout
+      - verbose: Verbosity flag
+      - authorization_context: Authorization context for downstream agents/tools
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        model: str | None = None,
+        verbose: bool | str | None = True,
+        timeout: int | None = 90,
+        authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
+        **_: Any,
+    ) -> None:
+        self.api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
+        self.api_base = (
+            api_base or os.environ.get("DATAROBOT_ENDPOINT") or "https://app.datarobot.com"
+        )
+        self.model = model
+        self.timeout = timeout if timeout is not None else 90
+        if isinstance(verbose, str):
+            self.verbose = verbose.lower() == "true"
+        elif verbose is None:
+            self.verbose = True
+        else:
+            self.verbose = bool(verbose)
+        self._mcp_tools: list[TTool] = []
+        self._authorization_context = authorization_context or {}
+        self._forwarded_headers: dict[str, str] = forwarded_headers or {}
+
+    def set_mcp_tools(self, tools: list[TTool]) -> None:
+        self._mcp_tools = tools
+
+    @property
+    def mcp_tools(self) -> list[TTool]:
+        """Return the list of MCP tools available to this agent.
+
+        Subclasses can use this to wire tools into CrewAI agents/tasks during
+        workflow construction inside ``build_crewai_workflow``.
+        """
+        return self._mcp_tools
+
+    @property
+    def authorization_context(self) -> dict[str, Any]:
+        """Return the authorization context for this agent."""
+        return self._authorization_context
+
+    @property
+    def forwarded_headers(self) -> dict[str, str]:
+        """Return the forwarded headers for this agent."""
+        return self._forwarded_headers
+
+    def litellm_api_base(self, deployment_id: str | None) -> str:
+        return get_api_base(self.api_base, deployment_id)
+
+    @abc.abstractmethod
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> "InvokeReturn":
+        raise NotImplementedError("Not implemented")
+
+    @classmethod
+    def create_pipeline_interactions_from_events(
+        cls,
+        events: list[Any] | None,
+    ) -> MultiTurnSample | None:
+        """Create a simple MultiTurnSample from a list of generic events/messages."""
+        if not events:
+            return None
+        return MultiTurnSample(user_input=events)
+
+
+def extract_user_prompt_content(
+    completion_create_params: CompletionCreateParams | Mapping[str, Any],
+) -> Any:
+    """Extract first user message content from OpenAI messages."""
+    params = cast(Mapping[str, Any], completion_create_params)
+    user_messages = [msg for msg in params.get("messages", []) if msg.get("role") == "user"]
+    # Get the last user message
+    user_prompt = user_messages[-1] if user_messages else {}
+    content = user_prompt.get("content", {})
+    # Try converting prompt from json to a dict
+    if isinstance(content, str):
+        try:
+            content = json.loads(content)
+        except json.JSONDecodeError:
+            pass
+
+    return content
+
+
+def make_system_prompt(suffix: str = "", *, prefix: str | None = None) -> str:
+    """Build a system prompt with optional prefix and suffix.
+
+    Parameters
+    ----------
+    suffix : str, default ""
+        Text appended after the prefix. If non-empty, it is placed on a new line.
+    prefix : str | None, keyword-only, default None
+        Custom prefix text. When ``None``, a default collaborative assistant
+        instruction is used.
+
+    Returns
+    -------
+    str
+        The composed system prompt string.
+    """
+    default_prefix = (
+        "You are a helpful AI assistant, collaborating with other assistants."
+        " Use the provided tools to progress towards answering the question."
+        " If you are unable to fully answer, that's OK, another assistant with different tools "
+        " will help where you left off. Execute what you can to make progress."
+    )
+    head = prefix if prefix is not None else default_prefix
+    if suffix:
+        return head + "\n" + suffix
+    return head
+
+
+# Structured type for token usage metrics in responses
+class UsageMetrics(TypedDict):
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+
+# Canonical return type for DRUM-compatible invoke implementations
+InvokeReturn = (
+    AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
+    | tuple[str, MultiTurnSample | None, UsageMetrics]
+)
+
+
+def default_usage_metrics() -> UsageMetrics:
+    """Return a metrics dict with required keys for OpenAI-compatible responses."""
+    return {
+        "completion_tokens": 0,
+        "prompt_tokens": 0,
+        "total_tokens": 0,
+    }
+
+
+def is_streaming(completion_create_params: CompletionCreateParams | Mapping[str, Any]) -> bool:
+    """Return True when the request asks for streaming, False otherwise.
+
+    Accepts both pydantic types and plain dictionaries.
+    """
+    params = cast(Mapping[str, Any], completion_create_params)
+    value = params.get("stream", False)
+    # Handle non-bool truthy values defensively (e.g., "true")
+    if isinstance(value, str):
+        return value.lower() == "true"
+    return bool(value)

datarobot_genai/core/chat/__init__.py

@@ -0,0 +1,19 @@
+"""Chat helpers and client utilities."""
+
+from .auth import initialize_authorization_context
+from .auth import resolve_authorization_context
+from .client import ToolClient
+from .responses import CustomModelChatResponse
+from .responses import CustomModelStreamingResponse
+from .responses import to_custom_model_chat_response
+from .responses import to_custom_model_streaming_response
+
+__all__ = [
+    "CustomModelChatResponse",
+    "CustomModelStreamingResponse",
+    "to_custom_model_chat_response",
+    "to_custom_model_streaming_response",
+    "ToolClient",
+    "resolve_authorization_context",
+    "initialize_authorization_context",
+]

datarobot_genai/core/chat/auth.py

@@ -0,0 +1,146 @@
+# 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.
+
+"""Authorization context helpers for chat flows."""
+
+from typing import Any
+
+from datarobot.models.genai.agent.auth import set_authorization_context
+from openai.types import CompletionCreateParams
+from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
+from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+
+def _get_authorization_context_from_headers(
+    headers: dict[str, str],
+    secret_key: str | None = None,
+) -> dict[str, Any] | None:
+    """Extract authorization context from headers using AuthContextHeaderHandler.
+
+    Parameters
+    ----------
+    headers : dict[str, str]
+        HTTP headers from which to extract the authorization context.
+    secret_key : str | None
+        Secret key for JWT decoding. If None, retrieves from environment variable.
+
+    Returns
+    -------
+    dict[str, Any] | None
+        The extracted authorization context, or None if not found.
+    """
+    handler = AuthContextHeaderHandler(secret_key=secret_key)
+    if context := handler.get_context(headers):
+        return context.model_dump()
+    return None
+
+
+def _get_authorization_context_from_params(
+    completion_create_params: CompletionCreateParams
+    | CompletionCreateParamsNonStreaming
+    | CompletionCreateParamsStreaming,
+) -> dict[str, Any] | None:
+    """Extract authorization context from completion create parameters.
+
+    Parameters
+    ----------
+    completion_create_params : CompletionCreateParams
+        The parameters used to create the completion.
+
+    Returns
+    -------
+    dict[str, Any] | None
+        The extracted authorization context, or None if not found.
+    """
+    return completion_create_params.get("authorization_context", None)
+
+
+def resolve_authorization_context(
+    completion_create_params: CompletionCreateParams
+    | CompletionCreateParamsNonStreaming
+    | CompletionCreateParamsStreaming,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Resolve the authorization context for the agent.
+
+    Authorization context is required for propagating information needed by downstream
+    agents and tools to retrieve access tokens to connect to external services. This method
+    extracts the authorization context from either the incoming HTTP headers or the completion
+    create parameters.
+
+    Parameters
+    ----------
+    completion_create_params : CompletionCreateParams | CompletionCreateParamsNonStreaming |
+        CompletionCreateParamsStreaming
+        Parameters supplied to the completion API. May include a fallback
+        ``authorization_context`` mapping under the same key.
+    **kwargs : Any
+        Additional keyword arguments. Expected to include a ``headers`` key
+        containing incoming HTTP headers as ``dict[str, str]``.
+
+    Returns
+    -------
+    dict[str, Any]
+        The initialized authorization context.
+    """
+    incoming_headers = kwargs.get("headers", {})
+
+    # Recommended way of propagating authorization context is via headers
+    # with JWT endoding/decoding for additional security. The completion params
+    # is used as a fallback for backward compatibility only and may be removed in
+    # the future.
+    authorization_context: dict[str, Any] = (
+        _get_authorization_context_from_headers(incoming_headers)
+        or _get_authorization_context_from_params(completion_create_params)
+        or {}
+    )
+
+    return authorization_context
+
+
+def initialize_authorization_context(
+    completion_create_params: CompletionCreateParams
+    | CompletionCreateParamsNonStreaming
+    | CompletionCreateParamsStreaming,
+    **kwargs: Any,
+) -> None:
+    """Set the authorization context for the agent.
+
+    Authorization context is required for propagating information needed by downstream
+    agents and tools to retrieve access tokens to connect to external services. When set,
+    authorization context will be automatically propagated when using ToolClient class.
+    authorization context will be propagated when using MCP Server component or when
+    using ToolClient class.
+
+    Parameters
+    ----------
+    completion_create_params : CompletionCreateParams | CompletionCreateParamsNonStreaming |
+        CompletionCreateParamsStreaming
+        Parameters supplied to the completion API. May include a fallback
+        ``authorization_context`` mapping under the same key.
+    **kwargs : Any
+        Additional keyword arguments. Expected to include a ``headers`` key
+        containing incoming HTTP headers as ``dict[str, str]``.
+
+    """
+    authorization_context = resolve_authorization_context(
+        completion_create_params,
+        **kwargs,
+    )
+
+    # Note: authorization context internally uses contextvars, which are
+    # thread-safe and async-safe.
+    set_authorization_context(authorization_context)

datarobot_genai/core/chat/client.py

@@ -0,0 +1,178 @@
+# 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.
+
+"""Client for interacting with Agent Tools deployments for chat and scoring."""
+
+import json
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import datarobot as dr
+import openai
+import pandas as pd
+from datarobot.models.genai.agent.auth import get_authorization_context
+from datarobot_predict.deployment import PredictionResult
+from datarobot_predict.deployment import UnstructuredPredictionResult
+from datarobot_predict.deployment import predict
+from datarobot_predict.deployment import predict_unstructured
+from openai.types import CompletionCreateParams
+from openai.types.chat import ChatCompletion
+from openai.types.chat import ChatCompletionChunk
+
+from ..utils.urls import get_api_base
+
+
+class ToolClient:
+    """Client for interacting with Agent Tools Deployments.
+
+    This class provides methods to call the custom model tool using various hooks:
+    `score`, `score_unstructured`, and `chat`. When the `authorization_context` is set,
+    the client automatically propagates it to the agent tool. The `authorization_context`
+    is required for retrieving access tokens to connect to external services.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        authorization_context: dict[str, Any] | None = None,
+    ):
+        """Initialize the ToolClient.
+
+        Args:
+            api_key (str | None): API key for authentication. Defaults to the
+                environment variable `DATAROBOT_API_TOKEN`.
+            base_url (str | None): Base URL for the DataRobot API. Defaults to the
+                environment variable `DATAROBOT_ENDPOINT` or 'app.datarobot.com'.
+            authorization_context (dict[str, Any] | None): Authorization context to use
+                for tool calls. If None, will attempt to get from ContextVar (for backward
+                compatibility).
+        """
+        self.api_key = api_key or os.getenv("DATAROBOT_API_TOKEN")
+        base_url = base_url or os.getenv("DATAROBOT_ENDPOINT") or "https://app.datarobot.com"
+        base_url = get_api_base(base_url, deployment_id=None)
+        self.base_url = base_url
+        self._authorization_context = authorization_context
+
+    @property
+    def datarobot_api_endpoint(self) -> str:
+        return self.base_url + "api/v2"
+
+    def get_deployment(self, deployment_id: str) -> dr.Deployment:
+        """Retrieve a deployment by its ID.
+
+        Args:
+            deployment_id (str): The ID of the deployment.
+
+        Returns
+        -------
+            dr.Deployment: The deployment object.
+        """
+        dr.Client(self.api_key, self.datarobot_api_endpoint)
+        return dr.Deployment.get(deployment_id=deployment_id)
+
+    def call(
+        self,
+        deployment_id: str,
+        payload: dict[str, Any],
+        authorization_context: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> UnstructuredPredictionResult:
+        """Run the custom model tool using score_unstructured hook.
+
+        Args:
+            deployment_id (str): The ID of the deployment.
+            payload (dict[str, Any]): The input payload.
+            authorization_context (dict[str, Any] | None): Authorization context to use.
+                If None, uses the context from initialization or falls back to ContextVar.
+            **kwargs: Additional keyword arguments.
+
+        Returns
+        -------
+            UnstructuredPredictionResult: The response content and headers.
+        """
+        # Use explicit context, fall back to instance context, then ContextVar
+        auth_ctx = authorization_context or self._authorization_context
+        if auth_ctx is None:
+            try:
+                auth_ctx = get_authorization_context()
+            except LookupError:
+                auth_ctx = {}
+
+        data = {
+            "payload": payload,
+            "authorization_context": auth_ctx,
+        }
+        return predict_unstructured(
+            deployment=self.get_deployment(deployment_id),
+            data=json.dumps(data),
+            content_type="application/json",
+            **kwargs,
+        )
+
+    def score(
+        self, deployment_id: str, data_frame: pd.DataFrame, **kwargs: Any
+    ) -> PredictionResult:
+        """Run the custom model tool using score hook.
+
+        Args:
+            deployment_id (str): The ID of the deployment.
+            data_frame (pd.DataFrame): The input data frame.
+            **kwargs: Additional keyword arguments.
+
+        Returns
+        -------
+            PredictionResult: The response content and headers.
+        """
+        return predict(
+            deployment=self.get_deployment(deployment_id),
+            data_frame=data_frame,
+            **kwargs,
+        )
+
+    def chat(
+        self,
+        completion_create_params: CompletionCreateParams,
+        model: str,
+        authorization_context: dict[str, Any] | None = None,
+    ) -> ChatCompletion | Iterator[ChatCompletionChunk]:
+        """Run the custom model tool with the chat hook.
+
+        Args:
+            completion_create_params (CompletionCreateParams): Parameters for the chat completion.
+            model (str): The model to use.
+            authorization_context (dict[str, Any] | None): Authorization context to use.
+                If None, uses the context from initialization or falls back to ContextVar.
+
+        Returns
+        -------
+            Union[ChatCompletion, Iterator[ChatCompletionChunk]]: The chat completion response.
+        """
+        # Use explicit context, fall back to instance context, then ContextVar
+        auth_ctx = authorization_context or self._authorization_context
+        if auth_ctx is None:
+            try:
+                auth_ctx = get_authorization_context()
+            except LookupError:
+                auth_ctx = {}
+
+        extra_body = {
+            "authorization_context": auth_ctx,
+        }
+        return openai.chat.completions.create(
+            **completion_create_params,
+            model=model,
+            extra_body=extra_body,
+        )