datarobot-genai 0.1.71__tar.gz → 0.2.5__tar.gz

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.71
+Version: 0.2.5
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -8,22 +8,25 @@ License: Apache-2.0
 License-File: AUTHORS
 License-File: LICENSE
 Requires-Python: <3.13,>=3.10
+Requires-Dist: ag-ui-protocol<0.2.0,>=0.1.9
 Requires-Dist: datarobot-drum<2.0.0,>=1.17.5
 Requires-Dist: datarobot-predict<2.0.0,>=1.13.2
-Requires-Dist: datarobot<4.0.0,>=3.9.1
+Requires-Dist: datarobot<4.0.0,>=3.10.0
 Requires-Dist: openai<2.0.0,>=1.76.2
 Requires-Dist: opentelemetry-instrumentation-aiohttp-client<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-httpx<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-openai<1.0.0,>=0.40.5
 Requires-Dist: opentelemetry-instrumentation-requests<1.0.0,>=0.43b0
+Requires-Dist: opentelemetry-instrumentation-threading<1.0.0,>=0.43b0
 Requires-Dist: pandas<3.0.0,>=2.2.3
+Requires-Dist: pyarrow==20.0.0
 Requires-Dist: pyjwt<3.0.0,>=2.10.1
 Requires-Dist: pypdf<7.0.0,>=6.1.3
 Requires-Dist: ragas<0.4.0,>=0.3.8
 Requires-Dist: requests<3.0.0,>=2.32.4
 Provides-Extra: crewai
 Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
-Requires-Dist: crewai<1.0.0,>=0.193.2; extra == 'crewai'
+Requires-Dist: crewai>=1.1.0; extra == 'crewai'
 Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'
 Requires-Dist: pybase64<2.0.0,>=1.4.2; extra == 'crewai'
 Provides-Extra: drmcp
@@ -56,8 +59,9 @@ Requires-Dist: llama-index<0.14.0,>=0.13.6; extra == 'llamaindex'
 Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'llamaindex'
 Requires-Dist: pypdf<7.0.0,>=6.0.0; extra == 'llamaindex'
 Provides-Extra: nat
+Requires-Dist: anyio==4.11.0; extra == 'nat'
+Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
-Requires-Dist: nvidia-nat-crewai==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.71"
+version = "0.2.5"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -12,7 +12,7 @@ license = { text = "Apache-2.0" }
 authors = [{ name = "DataRobot, Inc." }]
 dependencies = [
   "requests>=2.32.4,<3.0.0",
-  "datarobot>=3.9.1,<4.0.0",
+  "datarobot>=3.10.0,<4.0.0",
   "datarobot-predict>=1.13.2,<2.0.0",
   "openai>=1.76.2,<2.0.0",
   "pandas>=2.2.3,<3.0.0",
@@ -24,6 +24,9 @@ dependencies = [
   "opentelemetry-instrumentation-aiohttp-client>=0.43b0,<1.0.0",
   "opentelemetry-instrumentation-httpx>=0.43b0,<1.0.0",
   "opentelemetry-instrumentation-openai>=0.40.5,<1.0.0",
+  "opentelemetry-instrumentation-threading>=0.43b0,<1.0.0",
+  "ag-ui-protocol>=0.1.9,<0.2.0",
+  "pyarrow==20.0.0",
 ]
 
 [project.urls]
@@ -35,7 +38,7 @@ datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
 
 [project.optional-dependencies]
 crewai = [
-  "crewai>=0.193.2,<1.0.0",
+  "crewai>=1.1.0",
   "crewai-tools[mcp]>=0.69.0,<0.77.0",
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
   "pybase64>=1.4.2,<2.0.0",
@@ -59,12 +62,13 @@ llamaindex = [
 nat = [
   "nvidia-nat==1.3.0; python_version >= '3.11'",
   "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
-  "nvidia-nat-crewai==1.3.0; python_version >= '3.11'",
   "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
+  "crewai>=1.1.0; python_version >= '3.11'",
   "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Need this to support datarobot-llm plugin
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
   "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",
   "opentelemetry-instrumentation-langchain>=0.40.5,<1.0.0",
+  "anyio==4.11.0",
 ]
 pydanticai = [
   "logfire>=4.6.0,<5.0.0",
@@ -104,7 +108,7 @@ include = [
 [tool.uv]
 package = true
 override-dependencies = [
-  "ragas>=0.3.8",
+  "ragas>=0.3.8,<0.4.0",
 ]
 exclude-dependencies = [
   "uv",

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/core/agents/base.py

@@ -23,6 +23,7 @@ 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
 
@@ -167,7 +168,7 @@ class UsageMetrics(TypedDict):
 
 # Canonical return type for DRUM-compatible invoke implementations
 InvokeReturn = (
-    AsyncGenerator[tuple[str, MultiTurnSample | None, UsageMetrics], None]
+    AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
     | tuple[str, MultiTurnSample | None, UsageMetrics]
 )
 

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/core/chat/responses.py

@@ -14,14 +14,23 @@
 
 """OpenAI-compatible response helpers for chat interactions."""
 
+import asyncio
+import queue
 import time
 import traceback as tb
 import uuid
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator
+from collections.abc import AsyncIterator
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from typing import TypeVar
 
+from ag_ui.core import BaseEvent
+from ag_ui.core import Event
+from ag_ui.core import TextMessageChunkEvent
+from ag_ui.core import TextMessageContentEvent
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion
 from openai.types.chat import ChatCompletionChunk
@@ -40,6 +49,7 @@ class CustomModelChatResponse(ChatCompletion):
 
 class CustomModelStreamingResponse(ChatCompletionChunk):
     pipeline_interactions: str | None = None
+    event: Event | None = None
 
 
 def to_custom_model_chat_response(
@@ -83,7 +93,7 @@ def to_custom_model_streaming_response(
     thread_pool_executor: ThreadPoolExecutor,
     event_loop: AbstractEventLoop,
     streaming_response_generator: AsyncGenerator[
-        tuple[str, MultiTurnSample | None, dict[str, int]], None
+        tuple[str | Event, MultiTurnSample | None, dict[str, int]], None
     ],
     model: str | object | None,
 ) -> Iterator[CustomModelStreamingResponse]:
@@ -105,7 +115,7 @@ def to_custom_model_streaming_response(
         while True:
             try:
                 (
-                    response_text,
+                    response_text_or_event,
                     pipeline_interactions,
                     usage_metrics,
                 ) = thread_pool_executor.submit(
@@ -114,10 +124,10 @@ def to_custom_model_streaming_response(
                 last_pipeline_interactions = pipeline_interactions
                 last_usage_metrics = usage_metrics
 
-                if response_text:
+                if isinstance(response_text_or_event, str) and response_text_or_event:
                     choice = ChunkChoice(
                         index=0,
-                        delta=ChoiceDelta(role="assistant", content=response_text),
+                        delta=ChoiceDelta(role="assistant", content=response_text_or_event),
                         finish_reason=None,
                     )
                     yield CustomModelStreamingResponse(
@@ -130,6 +140,29 @@ def to_custom_model_streaming_response(
                         if usage_metrics
                         else None,
                     )
+                elif isinstance(response_text_or_event, BaseEvent):
+                    content = ""
+                    if isinstance(
+                        response_text_or_event, (TextMessageContentEvent, TextMessageChunkEvent)
+                    ):
+                        content = response_text_or_event.delta or content
+                    choice = ChunkChoice(
+                        index=0,
+                        delta=ChoiceDelta(role="assistant", content=content),
+                        finish_reason=None,
+                    )
+
+                    yield CustomModelStreamingResponse(
+                        id=completion_id,
+                        object="chat.completion.chunk",
+                        created=created,
+                        model=model,
+                        choices=[choice],
+                        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
+                        if usage_metrics
+                        else None,
+                        event=response_text_or_event,
+                    )
             except StopAsyncIteration:
                 break
         event_loop.run_until_complete(streaming_response_generator.aclose())
@@ -168,3 +201,97 @@ def to_custom_model_streaming_response(
             choices=[choice],
             usage=None,
         )
+
+
+def streaming_iterator_to_custom_model_streaming_response(
+    streaming_response_iterator: Iterator[tuple[str, MultiTurnSample | None, dict[str, int]]],
+    model: str | object | None,
+) -> Iterator[CustomModelStreamingResponse]:
+    """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
+    completion_id = str(uuid.uuid4())
+    created = int(time.time())
+
+    last_pipeline_interactions = None
+    last_usage_metrics = None
+
+    while True:
+        try:
+            (
+                response_text,
+                pipeline_interactions,
+                usage_metrics,
+            ) = next(streaming_response_iterator)
+            last_pipeline_interactions = pipeline_interactions
+            last_usage_metrics = usage_metrics
+
+            if response_text:
+                choice = ChunkChoice(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content=response_text),
+                    finish_reason=None,
+                )
+                yield CustomModelStreamingResponse(
+                    id=completion_id,
+                    object="chat.completion.chunk",
+                    created=created,
+                    model=model,
+                    choices=[choice],
+                    usage=CompletionUsage(**usage_metrics) if usage_metrics else None,
+                )
+        except StopIteration:
+            break
+    # Yield final chunk indicating end of stream
+    choice = ChunkChoice(
+        index=0,
+        delta=ChoiceDelta(role="assistant"),
+        finish_reason="stop",
+    )
+    yield CustomModelStreamingResponse(
+        id=completion_id,
+        object="chat.completion.chunk",
+        created=created,
+        model=model,
+        choices=[choice],
+        usage=CompletionUsage(**last_usage_metrics) if last_usage_metrics else None,
+        pipeline_interactions=last_pipeline_interactions.model_dump_json()
+        if last_pipeline_interactions
+        else None,
+    )
+
+
+T = TypeVar("T")
+
+
+def async_gen_to_sync_thread(
+    async_iterator: AsyncIterator[T],
+    thread_pool_executor: ThreadPoolExecutor,
+    event_loop: asyncio.AbstractEventLoop,
+) -> Iterator[T]:
+    """Run an async iterator in a separate thread and provide a sync iterator."""
+    # A thread-safe queue for communication
+    sync_queue: queue.Queue[Any] = queue.Queue()
+    # A sentinel object to signal the end of the async generator
+    SENTINEL = object()  # noqa: N806
+
+    async def run_async_to_queue() -> None:
+        """Run in the separate thread's event loop."""
+        try:
+            async for item in async_iterator:
+                sync_queue.put(item)
+        except Exception as e:
+            # Put the exception on the queue to be re-raised in the main thread
+            sync_queue.put(e)
+        finally:
+            # Signal the end of iteration
+            sync_queue.put(SENTINEL)
+
+    thread_pool_executor.submit(event_loop.run_until_complete, run_async_to_queue()).result()
+
+    # The main thread consumes items synchronously
+    while True:
+        item = sync_queue.get()
+        if item is SENTINEL:
+            break
+        if isinstance(item, Exception):
+            raise item
+        yield item

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/core/custom_model.py

@@ -26,7 +26,6 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from typing import Literal
 
-from datarobot_drum import RuntimeParameters
 from openai.types.chat import CompletionCreateParams
 from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
 from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
@@ -41,26 +40,6 @@ from datarobot_genai.core.telemetry_agent import instrument
 logger = logging.getLogger(__name__)
 
 
-def maybe_set_env_from_runtime_parameters(key: str) -> None:
-    """Set an environment variable from a DRUM Runtime Parameter if it exists.
-
-    This is safe to call outside of the DataRobot runtime. If the parameter is not available,
-    the function does nothing.
-    """
-    runtime_parameter_placeholder_value = "SET_VIA_PULUMI_OR_MANUALLY"
-    try:
-        runtime_parameter_value = RuntimeParameters.get(key)
-        if (
-            runtime_parameter_value
-            and len(runtime_parameter_value) > 0
-            and runtime_parameter_value != runtime_parameter_placeholder_value
-        ):
-            os.environ[key] = runtime_parameter_value
-    except ValueError:
-        # Local dev: runtime parameters may be unavailable
-        pass
-
-
 def load_model() -> tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]:
     """Initialize a dedicated event loop within a worker thread.
 
@@ -83,7 +62,6 @@ def chat_entrypoint(
     load_model_result: tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop],
     *,
     work_dir: str | None = None,
-    runtime_parameter_keys: list[str] | None = None,
     framework: Literal["crewai", "langgraph", "llamaindex", "nat"] | None = None,
     **kwargs: Any,
 ) -> CustomModelChatResponse | Iterator[CustomModelStreamingResponse]:
@@ -103,10 +81,6 @@ def chat_entrypoint(
     work_dir : Optional[str]
         Working directory to ``chdir`` into before invoking the agent. This is useful
         when relative paths are used in agent templates.
-    runtime_parameter_keys : Optional[List[str]]
-        Runtime parameter keys (DataRobot custom model) to propagate into env. When
-        ``None``, defaults to
-        ``['EXTERNAL_MCP_URL', 'MCP_DEPLOYMENT_ID']``.
     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.
@@ -129,12 +103,6 @@ def chat_entrypoint(
         except Exception as e:
             logger.warning(f"Failed to change working directory to {work_dir}: {e}")
 
-    # Load MCP runtime parameters and session secret if configured
-    if runtime_parameter_keys is None:
-        runtime_parameter_keys = ["EXTERNAL_MCP_URL", "MCP_DEPLOYMENT_ID"]
-    for key in runtime_parameter_keys:
-        maybe_set_env_from_runtime_parameters(key)
-
     # Retrieve authorization context using all supported methods for downstream agents/tools
     completion_create_params["authorization_context"] = resolve_authorization_context(
         completion_create_params, **kwargs

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/core/utils/auth.py

@@ -74,7 +74,22 @@ class AuthContextHeaderHandler:
         if algorithm is None:
             raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")
 
-        self.secret_key = secret_key or AuthContextConfig().session_secret_key
+        # 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
 

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/drmcp/core/config.py

@@ -197,6 +197,54 @@ class MCPServerConfig(BaseSettings):
         description="Enable/disable predictive tools",
     )
 
+    # Jira tools
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+    is_jira_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Jira OAuth provider is configured for Jira integration",
+    )
+
+    @property
+    def is_jira_oauth_configured(self) -> bool:
+        return self.is_jira_oauth_provider_configured or bool(
+            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
+        )
+
+    # Confluence tools
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+    is_confluence_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Confluence OAuth provider is configured for Confluence integration",
+    )
+
+    @property
+    def is_confluence_oauth_configured(self) -> bool:
+        return self.is_confluence_oauth_provider_configured or bool(
+            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
+        )
+
     @field_validator(
         "otel_attributes",
         mode="before",
@@ -220,6 +268,10 @@ class MCPServerConfig(BaseSettings):
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
         "enable_predictive_tools",
+        "enable_jira_tools",
+        "is_jira_oauth_provider_configured",
+        "enable_confluence_tools",
+        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -40,6 +40,8 @@ from .routes_utils import prefix_mount_path
 from .server_life_cycle import BaseServerLifecycle
 from .telemetry import OtelASGIMiddleware
 from .telemetry import initialize_telemetry
+from .tool_config import TOOL_CONFIGS
+from .tool_config import is_tool_enabled
 
 
 def _import_modules_from_dir(
@@ -115,6 +117,9 @@ class DataRobotMCPServer:
         self._mcp = mcp
         self._mcp_transport = transport
 
+        # Configure MCP server capabilities
+        self._configure_mcp_capabilities()
+
         # Initialize telemetry
         initialize_telemetry(mcp)
 
@@ -139,11 +144,12 @@ class DataRobotMCPServer:
 
         # Load static tools modules
         base_dir = os.path.dirname(os.path.dirname(__file__))
-        if self._config.enable_predictive_tools:
-            _import_modules_from_dir(
-                os.path.join(base_dir, "tools", "predictive"),
-                "datarobot_genai.drmcp.tools.predictive",
-            )
+        for tool_type, tool_config in TOOL_CONFIGS.items():
+            if is_tool_enabled(tool_type, self._config):
+                _import_modules_from_dir(
+                    os.path.join(base_dir, "tools", tool_config["directory"]),
+                    tool_config["package_prefix"],
+                )
 
         # Load memory management tools if available
         if self._memory_manager:
@@ -163,6 +169,37 @@ class DataRobotMCPServer:
         if transport == "streamable-http":
             register_routes(self._mcp)
 
+    def _configure_mcp_capabilities(self) -> None:
+        """Configure MCP capabilities that FastMCP doesn't expose directly.
+
+        See: https://github.com/modelcontextprotocol/python-sdk/issues/1126
+        """
+        server = self._mcp._mcp_server
+
+        # Declare prompts_changed capability  (capabilities.prompts.listChanged: true)
+        server.notification_options.prompts_changed = True
+
+        # Declare experimental capabilities ( experimental.dynamic_prompts: true)
+        server.experimental_capabilities = {"dynamic_prompts": {"enabled": True}}
+
+        # Patch to include experimental_capabilities (FastMCP doesn't expose this)
+        original = server.create_initialization_options
+
+        def patched(
+            notification_options: Any = None,
+            experimental_capabilities: dict[str, dict[str, Any]] | None = None,
+            **kwargs: Any,
+        ) -> Any:
+            if experimental_capabilities is None:
+                experimental_capabilities = getattr(server, "experimental_capabilities", None)
+            return original(
+                notification_options=notification_options,
+                experimental_capabilities=experimental_capabilities,
+                **kwargs,
+            )
+
+        server.create_initialization_options = patched
+
     def run(self, show_banner: bool = False) -> None:
         """Run the DataRobot MCP server synchronously."""
         try:
@@ -179,6 +216,9 @@ class DataRobotMCPServer:
                 self._logger.info("Registering dynamic prompts from prompt management...")
                 asyncio.run(register_prompts_from_datarobot_prompt_management())
 
+            # Execute pre-server start actions
+            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
+
             # List registered tools, prompts, and resources before starting server
             tools = asyncio.run(self._mcp._list_tools_mcp())
             prompts = asyncio.run(self._mcp._list_prompts_mcp())
@@ -198,9 +238,6 @@ class DataRobotMCPServer:
             for resource in resources:
                 self._logger.info(f" > {resource.name}")
 
-            # Execute pre-server start actions
-            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
-
             # Create event loop for async operations
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)

datarobot_genai-0.2.5/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py

@@ -0,0 +1,70 @@
+# 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 collections import defaultdict
+
+import datarobot as dr
+
+from datarobot_genai.drmcp.core.clients import get_api_client
+
+
+def get_datarobot_prompt_templates() -> list[dr.genai.PromptTemplate]:
+    try:
+        return dr.genai.PromptTemplate.list()
+    except Exception:
+        return []
+
+
+def get_datarobot_prompt_template_versions(
+    prompt_template_ids: list[str],
+) -> dict[str, list[dr.genai.PromptTemplateVersion]]:
+    # Still missing in SDK
+    prompt_template_versions_data = dr.utils.pagination.unpaginate(
+        initial_url="genai/promptTemplates/versions/",
+        initial_params={
+            "promptTemplateIds": prompt_template_ids,
+        },
+        client=get_api_client(),
+    )
+    prompt_template_versions = defaultdict(list)
+    for prompt_template_version in prompt_template_versions_data:
+        prompt_template_versions[prompt_template_version["promptTemplateId"]].append(
+            dr.genai.PromptTemplateVersion(
+                id=prompt_template_version["id"],
+                prompt_template_id=prompt_template_version["promptTemplateId"],
+                prompt_text=prompt_template_version["promptText"],
+                commit_comment=prompt_template_version["commitComment"],
+                version=prompt_template_version["version"],
+                variables=prompt_template_version["variables"],
+                creation_date=prompt_template_version["creationDate"],
+                creation_user_id=prompt_template_version["creationUserId"],
+                user_name=prompt_template_version["userName"],
+            )
+        )
+    return prompt_template_versions
+
+
+def get_datarobot_prompt_template(prompt_template_id: str) -> dr.genai.PromptTemplate | None:
+    try:
+        return dr.genai.PromptTemplate.get(prompt_template_id)
+    except Exception:
+        return None
+
+
+def get_datarobot_prompt_template_version(
+    prompt_template_id: str, prompt_template_version_id: str
+) -> dr.genai.PromptTemplateVersion | None:
+    try:
+        return dr.genai.PromptTemplateVersion.get(prompt_template_id, prompt_template_version_id)
+    except Exception:
+        return None

{datarobot_genai-0.1.71 → datarobot_genai-0.2.5}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -18,15 +18,13 @@ from collections.abc import Callable
 from inspect import Parameter
 from inspect import Signature
 
+import datarobot as dr
 from fastmcp.prompts.prompt import Prompt
 from pydantic import Field
 
 from datarobot_genai.drmcp.core.exceptions import DynamicPromptRegistrationError
 from datarobot_genai.drmcp.core.mcp_instance import register_prompt
 
-from .dr_lib import DrPrompt
-from .dr_lib import DrPromptVersion
-from .dr_lib import DrVariable
 from .dr_lib import get_datarobot_prompt_template_versions
 from .dr_lib import get_datarobot_prompt_templates
 
@@ -57,7 +55,8 @@ async def register_prompts_from_datarobot_prompt_management() -> None:
 
 
 async def register_prompt_from_datarobot_prompt_management(
-    prompt_template: DrPrompt, prompt_template_version: DrPromptVersion | None = None
+    prompt_template: dr.genai.PromptTemplate,
+    prompt_template_version: dr.genai.PromptTemplateVersion | None = None,
 ) -> Prompt:
     """Register a single prompt.
 
@@ -173,7 +172,7 @@ def to_valid_mcp_prompt_name(s: str) -> str:
 
 
 def make_prompt_function(
-    name: str, description: str, prompt_text: str, variables: list[DrVariable]
+    name: str, description: str, prompt_text: str, variables: list[dr.genai.Variable]
 ) -> Callable:
     params = []
     for v in variables: