datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/core/custom_model.py

@@ -0,0 +1,141 @@
+# 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.
+
+"""Helpers to implement DataRobot Custom Model chat entrypoints."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from collections.abc import AsyncGenerator
+from collections.abc import Iterator
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from typing import Literal
+
+from openai.types.chat import CompletionCreateParams
+from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
+from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
+
+from datarobot_genai.core.chat import CustomModelChatResponse
+from datarobot_genai.core.chat import CustomModelStreamingResponse
+from datarobot_genai.core.chat import to_custom_model_chat_response
+from datarobot_genai.core.chat import to_custom_model_streaming_response
+from datarobot_genai.core.chat.auth import resolve_authorization_context
+from datarobot_genai.core.telemetry_agent import instrument
+
+logger = logging.getLogger(__name__)
+
+
+def load_model() -> tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]:
+    """Initialize a dedicated event loop within a worker thread.
+
+    Returns
+    -------
+    (ThreadPoolExecutor, asyncio.AbstractEventLoop)
+        A single-worker executor and the associated event loop.
+    """
+    thread_pool_executor = ThreadPoolExecutor(1)
+    event_loop = asyncio.new_event_loop()
+    thread_pool_executor.submit(asyncio.set_event_loop, event_loop).result()
+    return (thread_pool_executor, event_loop)
+
+
+def chat_entrypoint(
+    agent_cls: type[Any],
+    completion_create_params: CompletionCreateParams
+    | CompletionCreateParamsNonStreaming
+    | CompletionCreateParamsStreaming,
+    load_model_result: tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop],
+    *,
+    work_dir: str | None = None,
+    framework: Literal["crewai", "langgraph", "llamaindex", "nat"] | None = None,
+    **kwargs: Any,
+) -> CustomModelChatResponse | Iterator[CustomModelStreamingResponse]:
+    """Run a generic Custom Model chat entrypoint for agent-based implementations.
+
+    Parameters
+    ----------
+    agent_cls : Type[Any]
+        The agent class to instantiate. Must define an ``async invoke(...)`` method
+        returning either:
+          - a tuple (response_text, pipeline_interactions, usage_metrics)
+          - or an async generator yielding (delta_text, pipeline_interactions, usage_metrics)
+    completion_create_params : CompletionCreateParams | ...
+        Parameters supplied by OpenAI-compatible Chat API.
+    load_model_result : tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]
+        Values returned by :func:`load_model`.
+    work_dir : Optional[str]
+        Working directory to ``chdir`` into before invoking the agent. This is useful
+        when relative paths are used in agent templates.
+    framework : Optional[Literal["crewai", "langgraph", "llamaindex", "nat"]]
+        When provided, idempotently instruments HTTP clients, OpenAI SDK, and the
+        given framework. If omitted, general instrumentation is still applied.
+    **kwargs : Any
+        Extra values forwarded for header-based auth context extraction.
+    """
+    thread_pool_executor, event_loop = load_model_result
+
+    # Set up telemetry (idempotent). When framework is provided, instrument it as well.
+    try:
+        instrument(framework)
+    except Exception:
+        # Instrumentation is best-effort; proceed regardless
+        pass
+
+    # Optionally change working directory for frameworks which rely on relative paths
+    if work_dir:
+        try:
+            os.chdir(work_dir)
+        except Exception as e:
+            logger.warning(f"Failed to change working directory to {work_dir}: {e}")
+
+    # Retrieve authorization context using all supported methods for downstream agents/tools
+    completion_create_params["authorization_context"] = resolve_authorization_context(
+        completion_create_params, **kwargs
+    )
+    # Keep only allowed headers from the forwarded_headers.
+    incoming_headers = kwargs.get("headers", {}) or {}
+    allowed_headers = {"x-datarobot-api-token", "x-datarobot-api-key"}
+    forwarded_headers = {k: v for k, v in incoming_headers.items() if k.lower() in allowed_headers}
+    completion_create_params["forwarded_headers"] = forwarded_headers
+
+    # Instantiate user agent with all supplied completion params including auth context
+    agent = agent_cls(**completion_create_params)
+
+    # Invoke the agent and check if it returns a generator or a tuple
+    result = thread_pool_executor.submit(
+        event_loop.run_until_complete,
+        agent.invoke(completion_create_params=completion_create_params),
+    ).result()
+
+    # Streaming response (async generator)
+    if isinstance(result, AsyncGenerator):
+        return to_custom_model_streaming_response(
+            thread_pool_executor,
+            event_loop,
+            result,
+            model=completion_create_params.get("model"),
+        )
+
+    # Non-streaming response
+    response_text, pipeline_interactions, usage_metrics = result
+    return to_custom_model_chat_response(
+        response_text,
+        pipeline_interactions,
+        usage_metrics,
+        model=completion_create_params.get("model"),
+    )

datarobot_genai/core/mcp/common.py

@@ -0,0 +1,218 @@
+# 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 json
+import logging
+import re
+from http import HTTPStatus
+from typing import Any
+from typing import Literal
+
+import requests
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from pydantic import field_validator
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+logger = logging.getLogger(__name__)
+
+
+class MCPConfig(DataRobotAppFrameworkBaseSettings):
+    """Configuration for MCP server connection.
+
+    Derived values are exposed as properties rather than stored, avoiding
+    Pydantic field validation/serialization concerns for internal helpers.
+    """
+
+    external_mcp_url: str | None = None
+    external_mcp_headers: str | None = None
+    external_mcp_transport: Literal["sse", "streamable-http"] = "streamable-http"
+    mcp_deployment_id: str | None = None
+    datarobot_endpoint: str | None = None
+    datarobot_api_token: str | None = None
+    authorization_context: dict[str, Any] | None = None
+    forwarded_headers: dict[str, str] | None = None
+    mcp_server_port: int | None = None
+
+    _auth_context_handler: AuthContextHeaderHandler | None = None
+    _server_config: dict[str, Any] | None = None
+
+    @field_validator("external_mcp_headers", mode="before")
+    @classmethod
+    def validate_external_mcp_headers(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        candidate = value.strip()
+
+        try:
+            json.loads(candidate)
+        except json.JSONDecodeError:
+            msg = "external_mcp_headers must be valid JSON"
+            logger.warning(msg)
+            return None
+
+        return candidate
+
+    @field_validator("mcp_deployment_id", mode="before")
+    @classmethod
+    def validate_mcp_deployment_id(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        candidate = value.strip()
+
+        if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
+            msg = "mcp_deployment_id must be a valid 24-character hex ID"
+            logger.warning(msg)
+            return None
+
+        return candidate
+
+    def _authorization_bearer_header(self) -> dict[str, str]:
+        """Return Authorization header with Bearer token or empty dict."""
+        if not self.datarobot_api_token:
+            return {}
+        auth = (
+            self.datarobot_api_token
+            if self.datarobot_api_token.startswith("Bearer ")
+            else f"Bearer {self.datarobot_api_token}"
+        )
+        return {"Authorization": auth}
+
+    @property
+    def auth_context_handler(self) -> AuthContextHeaderHandler:
+        if self._auth_context_handler is None:
+            self._auth_context_handler = AuthContextHeaderHandler()
+        return self._auth_context_handler
+
+    @property
+    def server_config(self) -> dict[str, Any] | None:
+        if self._server_config is None:
+            self._server_config = self._build_server_config()
+        return self._server_config
+
+    def _authorization_context_header(self) -> dict[str, str]:
+        """Return X-DataRobot-Authorization-Context header or empty dict."""
+        try:
+            return self.auth_context_handler.get_header(self.authorization_context)
+        except (LookupError, RuntimeError):
+            # Authorization context not available (e.g., in tests)
+            return {}
+
+    def _build_authenticated_headers(self) -> dict[str, str]:
+        """Build headers for authenticated requests.
+
+        Returns
+        -------
+            Dictionary containing forwarded headers (if available) and authentication headers.
+        """
+        headers: dict[str, str] = {}
+        if self.forwarded_headers:
+            headers.update(self.forwarded_headers)
+        headers.update(self._authorization_bearer_header())
+        headers.update(self._authorization_context_header())
+        return headers
+
+    def _check_localhost_server(self, url: str, timeout: float = 2.0) -> bool:
+        """Check if MCP server is running on localhost.
+
+        Parameters
+        ----------
+        url : str
+            The URL to check.
+        timeout : float, optional
+            Request timeout in seconds (default: 2.0).
+
+        Returns
+        -------
+        bool
+            True if server is running and responding with OK status, False otherwise.
+        """
+        try:
+            response = requests.get(url, timeout=timeout)
+            return (
+                response.status_code == HTTPStatus.OK
+                and response.json().get("message") == "DataRobot MCP Server is running"
+            )
+        except requests.RequestException as e:
+            logger.debug(f"Failed to connect to MCP server at {url}: {e}")
+            return False
+
+    def _build_server_config(self) -> dict[str, Any] | None:
+        """
+        Get MCP server configuration.
+
+        Returns
+        -------
+            Server configuration dict with url, transport, and optional headers,
+            or None if not configured.
+        """
+        if self.mcp_deployment_id:
+            # DataRobot deployment ID - requires authentication
+            if self.datarobot_endpoint is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_endpoint must be set."
+                )
+            if self.datarobot_api_token is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_api_token must be set."
+                )
+
+            base_url = self.datarobot_endpoint.rstrip("/")
+            if not base_url.endswith("/api/v2"):
+                base_url = f"{base_url}/api/v2"
+
+            url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+            headers = self._build_authenticated_headers()
+
+            logger.info(f"Using DataRobot hosted MCP deployment: {url}")
+
+            return {
+                "url": url,
+                "transport": "streamable-http",
+                "headers": headers,
+            }
+
+        if self.external_mcp_url:
+            # External MCP URL - no authentication needed
+            headers = {}
+
+            # Merge external headers if provided
+            if self.external_mcp_headers:
+                external_headers = json.loads(self.external_mcp_headers)
+                headers.update(external_headers)
+
+            logger.info(f"Using external MCP URL: {self.external_mcp_url}")
+
+            return {
+                "url": self.external_mcp_url.rstrip("/"),
+                "transport": self.external_mcp_transport,
+                "headers": headers,
+            }
+
+        # No MCP configuration found, setup localhost if running locally
+        if self.mcp_server_port:
+            url = f"http://localhost:{self.mcp_server_port}"
+            if self._check_localhost_server(url):
+                headers = self._build_authenticated_headers()
+                logger.info(f"Using localhost MCP server: {url}")
+                return {
+                    "url": f"{url}/mcp",
+                    "transport": "streamable-http",
+                    "headers": headers,
+                }
+            logger.warning(f"MCP server is not running or not responding at {url}")
+
+        return None

datarobot_genai/core/telemetry_agent.py

@@ -0,0 +1,126 @@
+# 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.
+
+"""Lightweight, idempotent client/framework instrumentation for agents."""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+from typing import Any
+from typing import Literal
+from typing import cast
+
+# Suppress the "Attempting to instrument while already instrumented" warning
+logging.getLogger("opentelemetry.instrumentation.instrumentor").setLevel(logging.ERROR)
+logger = logging.getLogger(__name__)
+
+# Internal instrumentation state to avoid 'global' mutation warnings
+_INSTRUMENTATION_STATE = {"http": False, "openai": False, "threading": False}
+_INSTRUMENTED_FRAMEWORKS: set[str] = set()
+
+
+def _instrument_threading() -> None:
+    if _INSTRUMENTATION_STATE["threading"]:
+        return
+    try:
+        threading_module = importlib.import_module("opentelemetry.instrumentation.threading")
+        threading_instrumentor = getattr(threading_module, "ThreadingInstrumentor")
+        threading_instrumentor().instrument()
+        _INSTRUMENTATION_STATE["threading"] = True
+    except Exception as e:
+        logger.debug(f"threading instrumentation skipped: {e}")
+
+
+def _instrument_http_clients() -> None:
+    if _INSTRUMENTATION_STATE["http"]:
+        return
+    try:
+        requests_module = importlib.import_module("opentelemetry.instrumentation.requests")
+        requests_instrumentor = getattr(requests_module, "RequestsInstrumentor")
+        requests_instrumentor().instrument()
+    except Exception as e:
+        logger.debug(f"requests instrumentation skipped: {e}")
+    try:
+        aiohttp_module = importlib.import_module("opentelemetry.instrumentation.aiohttp_client")
+        aiohttp_instrumentor = getattr(aiohttp_module, "AioHttpClientInstrumentor")
+        aiohttp_instrumentor().instrument()
+    except Exception as e:
+        logger.debug(f"aiohttp instrumentation skipped: {e}")
+    try:
+        httpx_module = importlib.import_module("opentelemetry.instrumentation.httpx")
+        httpx_instrumentor = getattr(httpx_module, "HTTPXClientInstrumentor")
+        httpx_instrumentor().instrument()
+    except Exception as e:
+        logger.debug(f"httpx instrumentation skipped: {e}")
+    _INSTRUMENTATION_STATE["http"] = True
+
+
+def _instrument_openai() -> None:
+    if _INSTRUMENTATION_STATE["openai"]:
+        return
+    try:
+        openai_module = importlib.import_module("opentelemetry.instrumentation.openai")
+        openai_instrumentor = getattr(openai_module, "OpenAIInstrumentor")
+        openai_instrumentor().instrument()
+        _INSTRUMENTATION_STATE["openai"] = True
+    except Exception as e:
+        logger.debug(f"openai instrumentation skipped: {e}")
+
+
+def _instrument_framework(framework: str) -> None:
+    if framework in _INSTRUMENTED_FRAMEWORKS:
+        return
+    try:
+        if framework == "crewai":
+            crewai_module = importlib.import_module("opentelemetry.instrumentation.crewai")
+            crewai_instrumentor = getattr(crewai_module, "CrewAIInstrumentor")
+            crewai_instrumentor().instrument()
+            os.environ.setdefault("CREWAI_TESTING", "true")
+        elif framework == "langgraph":
+            # Provided by opentelemetry-instrumentation-langchain
+            langchain_module = importlib.import_module("opentelemetry.instrumentation.langchain")
+            langchain_instrumentor = getattr(langchain_module, "LangchainInstrumentor")
+            langchain_instrumentor().instrument()
+        elif framework == "llamaindex":
+            llamaindex_module = importlib.import_module("opentelemetry.instrumentation.llamaindex")
+            llamaindex_instrumentor = getattr(llamaindex_module, "LlamaIndexInstrumentor")
+            # LlamaIndex instrumentor lacks precise typing; cast to Any to avoid mypy complaints
+            cast(Any, llamaindex_instrumentor()).instrument()
+        elif framework == "nat":
+            _instrument_framework("crewai")
+            _instrument_framework("langgraph")
+            _instrument_framework("llamaindex")
+        _INSTRUMENTED_FRAMEWORKS.add(framework)
+    except Exception as e:
+        logger.debug(f"{framework} instrumentation skipped: {e}")
+
+
+def instrument(
+    framework: Literal["crewai", "langgraph", "llamaindex", "nat"] | None = None,
+) -> None:
+    """Idempotently instrument supported HTTP clients, OpenAI SDK, and optionally a framework.
+
+    Also disables telemetry for some third-party libraries to avoid duplicate/undesired tracking.
+    """
+    # Some libraries collect telemetry data by default. Disable that.
+    os.environ.setdefault("RAGAS_DO_NOT_TRACK", "true")
+    os.environ.setdefault("DEEPEVAL_TELEMETRY_OPT_OUT", "YES")
+
+    _instrument_threading()
+    _instrument_http_clients()
+    _instrument_openai()
+    if framework:
+        _instrument_framework(framework)

datarobot_genai/core/utils/__init__.py

@@ -0,0 +1,3 @@
+from .urls import get_api_base
+
+__all__ = ["get_api_base"]