datarobot-genai 0.1.63__tar.gz → 0.1.65__tar.gz

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/PKG-INFO

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

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/pyproject.toml

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

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/src/datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -184,13 +184,17 @@ class DataRobotMCPServer:
             prompts = asyncio.run(self._mcp._list_prompts_mcp())
             resources = asyncio.run(self._mcp._list_resources_mcp())
 
-            self._logger.info(f"Registered tools: {len(tools)}")
+            tools_count = len(tools)
+            prompts_count = len(prompts)
+            resources_count = len(resources)
+
+            self._logger.info(f"Registered tools: {tools_count}")
             for tool in tools:
                 self._logger.info(f" > {tool.name}")
-            self._logger.info(f"Registered prompts: {len(prompts)}")
+            self._logger.info(f"Registered prompts: {prompts_count}")
             for prompt in prompts:
                 self._logger.info(f" > {prompt.name}")
-            self._logger.info(f"Registered resources: {len(resources)}")
+            self._logger.info(f"Registered resources: {resources_count}")
             for resource in resources:
                 self._logger.info(f" > {resource.name}")
 
@@ -209,6 +213,9 @@ class DataRobotMCPServer:
                         self._mcp,
                         self._mcp_transport,
                         port=self._config.mcp_server_port,
+                        tools_count=tools_count,
+                        prompts_count=prompts_count,
+                        resources_count=resources_count,
                     )
 
                 if self._mcp_transport == "stdio":

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py

@@ -54,6 +54,9 @@ def log_server_custom_banner(
     host: str | None = None,
     port: int | None = None,
     path: str | None = None,
+    tools_count: int | None = None,
+    prompts_count: int | None = None,
+    resources_count: int | None = None,
 ) -> None:
     """
     Create and log a formatted banner with server information and logo.
@@ -64,13 +67,20 @@ def log_server_custom_banner(
         host: Host address (for HTTP transports)
         port: Port number (for HTTP transports)
         path: Server path (for HTTP transports)
+        tools_count: Number of tools registered
+        prompts_count: Number of prompts registered
+        resources_count: Number of resources registered
     """
     # Create the logo text
     # Use Text with no_wrap and markup disabled to preserve ANSI escape codes
     logo_text = Text.from_ansi(DR_LOGO_ASCII, no_wrap=True)
 
     # Create the main title
-    title_text = Text(f"DataRobot MCP Server {datarobot_genai_version}", style="bold green")
+    title_text = Text(f"DataRobot MCP Server {datarobot_genai_version}", style="dim green")
+    stats_text = Text(
+        f"{tools_count} tools, {prompts_count} prompts, {resources_count} resources",
+        style="bold green",
+    )
 
     # Create the information table
     info_table = Table.grid(padding=(0, 1))
@@ -107,6 +117,7 @@ def log_server_custom_banner(
         Align.center(logo_text),
         "",
         Align.center(title_text),
+        Align.center(stats_text),
         "",
         "",
         Align.center(info_table),

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py

@@ -11,7 +11,7 @@
 # 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
@@ -34,6 +34,7 @@ class DrVariable:
 @dataclass
 class DrPromptVersion:
     id: str
+    prompt_template_id: str
     version: int
     prompt_text: str
     variables: list[DrVariable]
@@ -45,6 +46,7 @@ class DrPromptVersion:
         ]
         return cls(
             id=d["id"],
+            prompt_template_id=d["promptTemplateId"],
             version=d["version"],
             prompt_text=d["promptText"],
             variables=variables,
@@ -58,7 +60,9 @@ class DrPrompt:
     description: str
 
     def get_latest_version(self) -> DrPromptVersion | None:
-        prompt_template_versions = get_datarobot_prompt_template_versions(self.id)
+        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)
@@ -77,15 +81,21 @@ def get_datarobot_prompt_templates() -> list[DrPrompt]:
     return [DrPrompt.from_dict(prompt_template) for prompt_template in prompt_templates_data]
 
 
-def get_datarobot_prompt_template_versions(prompt_template_id: str) -> list[DrPromptVersion]:
+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=f"genai/promptTemplates/{prompt_template_id}/versions/",
-        initial_params={},
+        initial_url="genai/promptTemplates/versions/",
+        initial_params={
+            "promptTemplateIds": prompt_template_ids,
+        },
         client=get_api_client(),
     )
-    prompt_template_versions = []
+    prompt_template_versions = defaultdict(list)
     for prompt_template_version in prompt_template_versions_data:
-        prompt_template_versions.append(DrPromptVersion.from_dict(prompt_template_version))
+        prompt_template_versions[prompt_template_version["promptTemplateId"]].append(
+            DrPromptVersion.from_dict(prompt_template_version)
+        )
     return prompt_template_versions
 
 

{datarobot_genai-0.1.63 → datarobot_genai-0.1.65}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -27,6 +27,7 @@ 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
 
 logger = logging.getLogger(__name__)
@@ -36,11 +37,21 @@ async def register_prompts_from_datarobot_prompt_management() -> None:
     """Register prompts from DataRobot Prompt Management."""
     prompts = get_datarobot_prompt_templates()
     logger.info(f"Found {len(prompts)} prompts in Prompts Management.")
+    all_prompts_versions = get_datarobot_prompt_template_versions(
+        prompt_template_ids=list({prompt.id for prompt in prompts})
+    )
 
     # Try to register each prompt, continue on failure
     for prompt in prompts:
+        prompt_versions = all_prompts_versions.get(prompt.id)
+        if not prompt_versions:
+            logger.warning(f"Prompt template id {prompt.id} has no versions.")
+            continue
+
+        latest_version = max(prompt_versions, key=lambda v: v.version)
+
         try:
-            await register_prompt_from_datarobot_prompt_management(prompt)
+            await register_prompt_from_datarobot_prompt_management(prompt, latest_version)
         except DynamicPromptRegistrationError:
             pass
 
@@ -114,15 +125,36 @@ async def register_prompt_from_datarobot_prompt_management(
         ) from exc
 
 
+def _escape_non_ascii(s: str) -> str:
+    out = []
+    for ch in s:
+        # If its space -> change to underscore
+        if ch.isspace():
+            out.append("_")
+        # ASCII letter, digit or underscore -> keep
+        elif ch.isascii() and (ch.isalnum() or ch == "_"):
+            out.append(ch)
+        # Everything else -> encode as 'xHEX'
+        else:
+            out.append(f"x{ord(ch):x}")
+    return "".join(out)
+
+
 def to_valid_mcp_prompt_name(s: str) -> str:
     """Convert an arbitrary string into a valid MCP prompt name."""
-    # Replace any sequence of invalid characters with '_'
-    s = re.sub(r"[^0-9a-zA-Z_]+", "_", s)
-
     # If its ONLY numbers return "prompt_[number]"
     if s.isdigit():
         return f"prompt_{s}"
 
+    # First, ASCII-transliterate using hex escape for non-ASCII
+    if not s.isascii():
+        # whole string non-ascii? -> escape and prefix with prompt_
+        encoded = _escape_non_ascii(s)
+        return f"prompt_{encoded}"
+
+    # Replace any sequence of invalid characters with '_'
+    s = re.sub(r"[^0-9a-zA-Z_]+", "_", s)
+
     # Remove leading characters that are not letters or underscores (can't start with a digit or _)
     s = re.sub(r"^[^a-zA-Z]+", "", s)