datarobot-genai 0.2.0__tar.gz → 0.2.1__tar.gz

{datarobot_genai-0.2.0 → datarobot_genai-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.0
+Version: 0.2.1
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -11,7 +11,7 @@ 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

{datarobot_genai-0.2.0 → datarobot_genai-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.0"
+version = "0.2.1"
 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",

datarobot_genai-0.2.1/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.2.0 → datarobot_genai-0.2.1}/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:

{datarobot_genai-0.2.0 → datarobot_genai-0.2.1}/src/datarobot_genai/drmcp/core/mcp_instance.py

@@ -286,6 +286,9 @@ class TaggedFastMCP(FastMCP):
                 f"already mapped to {existing_prompt_template_version_id}. "
                 f"Updating to version id = {prompt_template_version_id} and name = {prompt_name}"
             )
+            await self.remove_prompt_mapping(
+                prompt_template_id, existing_prompt_template_version_id
+            )
 
         self._prompts_map[prompt_template_id] = (prompt_template_version_id, prompt_name)
 
@@ -526,13 +529,14 @@ async def register_prompt(
     )
 
     # Register the prompt
-    registered_prompt = mcp.add_prompt(prompt)
     if prompt_template:
         prompt_template_id, prompt_template_version_id = prompt_template
         await mcp.set_prompt_mapping(
             prompt_template_id, prompt_template_version_id, prompt_name_no_duplicate
         )
 
+    registered_prompt = mcp.add_prompt(prompt)
+
     # Verify prompt is registered
     prompts = await mcp.get_prompts()
     if not any(prompt.name == prompt_name_no_duplicate for prompt in prompts.values()):

{datarobot_genai-0.2.0 → datarobot_genai-0.2.1}/src/datarobot_genai/drmcp/core/routes.py

@@ -428,7 +428,10 @@ def register_routes(mcp: TaggedFastMCP) -> None:
         """Refresh prompt templates."""
         try:
             await refresh_registered_prompt_template()
-            return JSONResponse(status_code=HTTPStatus.NO_CONTENT, content=None)
+            return JSONResponse(
+                status_code=HTTPStatus.OK,
+                content={"message": "Prompts refreshed successfully"},
+            )
         except Exception as e:
             return JSONResponse(
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR,

{datarobot_genai-0.2.0 → datarobot_genai-0.2.1}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -16,6 +16,8 @@ import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 
+import aiohttp
+from aiohttp import ClientSession as HttpClientSession
 from mcp import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
 
@@ -29,6 +31,11 @@ def get_dr_mcp_server_url() -> str | None:
     return os.environ.get("DR_MCP_SERVER_URL")
 
 
+def get_dr_mcp_server_http_url() -> str | None:
+    """Get DataRobot MCP server http URL."""
+    return os.environ.get("DR_MCP_SERVER_HTTP_URL")
+
+
 def get_openai_llm_client_config() -> dict[str, str]:
     """Get OpenAI LLM client configuration."""
     openai_api_key = os.environ.get("OPENAI_API_KEY")
@@ -94,3 +101,25 @@ async def ete_test_mcp_session(
                 yield session
     except asyncio.TimeoutError:
         raise TimeoutError(f"Check if the MCP server is running at {get_dr_mcp_server_url()}")
+
+
+@asynccontextmanager
+async def ete_test_http_session(
+    additional_headers: dict[str, str] | None = None,
+) -> AsyncGenerator[HttpClientSession, None]:
+    """Create an HTTP session for each test that can connect to MCP custom http routes.
+
+    Parameters
+    ----------
+    additional_headers : dict[str, str], optional
+        Additional headers to include in the HTTP session (e.g., auth headers for testing).
+    """
+    headers = get_headers()
+    if additional_headers:
+        headers.update(additional_headers)
+
+    async with ete_test_mcp_session(additional_headers=additional_headers):
+        async with aiohttp.ClientSession(
+            base_url=get_dr_mcp_server_http_url(), headers=headers
+        ) as client:
+            yield client

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

@@ -1,128 +0,0 @@
-# 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
-from dataclasses import dataclass
-
-import datarobot as dr
-
-from datarobot_genai.drmcp.core.clients import get_api_client
-
-# Needed SDK version (3.10.0) is not published yet. We'll reimplement simplified version of it.
-# get_datarobot_prompt_templates = dr.genai.PromptTemplate.list()
-# DrPrompt = dr.genai.PromptTemplate
-# DrPromptVersion = dr.genai.PromptTemplateVersion
-# DrVariable = dr.genai.Variable
-
-
-@dataclass
-class DrVariable:
-    name: str
-    description: str
-
-
-@dataclass
-class DrPromptVersion:
-    id: str
-    prompt_template_id: str
-    version: int
-    prompt_text: str
-    variables: list[DrVariable]
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "DrPromptVersion":
-        variables = [
-            DrVariable(name=v["name"], description=v["description"]) for v in d["variables"]
-        ]
-        return cls(
-            id=d["id"],
-            prompt_template_id=d["promptTemplateId"],
-            version=d["version"],
-            prompt_text=d["promptText"],
-            variables=variables,
-        )
-
-
-@dataclass
-class DrPrompt:
-    id: str
-    name: str
-    description: str
-
-    def get_latest_version(self) -> DrPromptVersion | None:
-        all_prompt_template_versions = get_datarobot_prompt_template_versions([self.id])
-        prompt_template_versions = all_prompt_template_versions.get(self.id)
-
-        if not prompt_template_versions:
-            return None
-        latest_version = max(prompt_template_versions, key=lambda v: v.version)
-        return latest_version
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "DrPrompt":
-        return cls(id=d["id"], name=d["name"], description=d["description"])
-
-
-def get_datarobot_prompt_templates() -> list[DrPrompt]:
-    prompt_templates_data = dr.utils.pagination.unpaginate(
-        initial_url="genai/promptTemplates/", initial_params={}, client=get_api_client()
-    )
-
-    return [DrPrompt.from_dict(prompt_template) for prompt_template in prompt_templates_data]
-
-
-def get_datarobot_prompt_template_versions(
-    prompt_template_ids: list[str],
-) -> dict[str, list[DrPromptVersion]]:
-    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(
-            DrPromptVersion.from_dict(prompt_template_version)
-        )
-    return prompt_template_versions
-
-
-def get_datarobot_prompt_template(prompt_template_id: str) -> DrPrompt | None:
-    api_client = get_api_client()
-    try:
-        prompt_template_response = api_client.get(
-            f"genai/promptTemplates/{prompt_template_id}/", join_endpoint=True
-        )
-        prompt_template_json = prompt_template_response.json()
-    except Exception:
-        return None
-
-    return DrPrompt.from_dict(prompt_template_json)
-
-
-def get_datarobot_prompt_template_version(
-    prompt_template_id: str, prompt_template_version_id: str
-) -> DrPromptVersion | None:
-    api_client = get_api_client()
-    try:
-        prompt_template_version_response = api_client.get(
-            f"genai/promptTemplates/{prompt_template_id}/versions/{prompt_template_version_id}/",
-            join_endpoint=True,
-        )
-        prompt_template_version_json = prompt_template_version_response.json()
-    except Exception:
-        return None
-
-    return DrPromptVersion.from_dict(prompt_template_version_json)