datarobot-genai 0.1.67__tar.gz → 0.1.68__tar.gz

{datarobot_genai-0.1.67 → datarobot_genai-0.1.68}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.67
+Version: 0.1.68
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.1.67 → datarobot_genai-0.1.68}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.67"
+version = "0.1.68"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.1.67 → datarobot_genai-0.1.68}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py

@@ -18,6 +18,8 @@ from fastmcp.prompts.prompt import Prompt
 
 from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template
 from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_version
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_versions
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_templates
 from datarobot_genai.drmcp.core.dynamic_prompts.register import (
     register_prompt_from_datarobot_prompt_management,
 )
@@ -83,3 +85,46 @@ async def delete_registered_prompt_template(prompt_template_id: str) -> bool:
         f"version {prompt_template_version_id}"
     )
     return True
+
+
+async def refresh_registered_prompt_template() -> None:
+    """Refresh all registered prompt templates in the MCP instance."""
+    prompt_templates = get_datarobot_prompt_templates()
+    prompt_templates_ids = {p.id for p in prompt_templates}
+    prompt_templates_versions = get_datarobot_prompt_template_versions(list(prompt_templates_ids))
+
+    mcp_prompt_templates_mappings = await mcp.get_prompt_mapping()
+
+    for prompt_template in prompt_templates:
+        prompt_template_versions = prompt_templates_versions.get(prompt_template.id)
+        if not prompt_template_versions:
+            continue
+
+        latest_version = max(prompt_template_versions, key=lambda v: v.version)
+
+        if prompt_template.id not in mcp_prompt_templates_mappings:
+            # New prompt template -> add
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        mcp_prompt_template_version, mcp_prompt = mcp_prompt_templates_mappings[prompt_template.id]
+
+        if mcp_prompt_template_version != latest_version:
+            # Current version saved in MCP is not the latest one => update it
+            await register_prompt_from_datarobot_prompt_management(
+                prompt_template=prompt_template, prompt_template_version=latest_version
+            )
+            continue
+
+        # Else => mcp_prompt_template_version == latest_version
+        # For now it means nothing changed as there's no possibility to edit promp template version.
+
+    for mcp_prompt_template_id, (
+        mcp_prompt_template_version_id,
+        _,
+    ) in mcp_prompt_templates_mappings.items():
+        if mcp_prompt_template_id not in prompt_templates_ids:
+            # We need to also delete prompt templates that are
+            await mcp.remove_prompt_mapping(mcp_prompt_template_id, mcp_prompt_template_version_id)

{datarobot_genai-0.1.67 → datarobot_genai-0.1.68}/src/datarobot_genai/drmcp/core/routes.py

@@ -19,6 +19,7 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse
 
 from .dynamic_prompts.controllers import delete_registered_prompt_template
+from .dynamic_prompts.controllers import refresh_registered_prompt_template
 from .dynamic_prompts.controllers import register_prompt_from_prompt_template_id_and_version
 from .dynamic_tools.deployment.controllers import delete_registered_tool_deployment
 from .dynamic_tools.deployment.controllers import get_registered_tool_deployments
@@ -418,6 +419,18 @@ def register_routes(mcp: TaggedFastMCP) -> None:
             )
         except Exception as e:
             return JSONResponse(
-                status_code=HTTPStatus.BAD_REQUEST,
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 content={"error": f"Failed to add prompt template: {str(e)}"},
             )
+
+    @mcp.custom_route(prefix_mount_path("/registeredPrompts"), methods=["PUT"])
+    async def refresh_prompt_templates(_: Request) -> JSONResponse:
+        """Refresh prompt templates."""
+        try:
+            await refresh_registered_prompt_template()
+            return JSONResponse(status_code=HTTPStatus.NO_CONTENT, content=None)
+        except Exception as e:
+            return JSONResponse(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                content={"error": f"Failed to refresh prompt templates: {str(e)}"},
+            )