datarobot-genai 0.1.62__py3-none-any.whl → 0.1.64__py3-none-any.whl

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/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)
 

datarobot_genai/langgraph/agent.py

@@ -32,6 +32,7 @@ from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
 from datarobot_genai.core.agents.base import UsageMetrics
 from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
 from datarobot_genai.langgraph.mcp import mcp_tools_context
 
 logger = logging.getLogger(__name__)
@@ -136,10 +137,10 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
         # The following code demonstrate both a synchronous and streaming response.
         # You can choose one or the other based on your use case, they function the same.
         # The main difference is returning a generator for streaming or a final response for sync.
-        if completion_create_params.get("stream"):
+        if is_streaming(completion_create_params):
             # Streaming response: yield each message as it is generated
             async def stream_generator() -> AsyncGenerator[
-                tuple[str, Any | None, UsageMetrics], None
+                tuple[str, MultiTurnSample | None, UsageMetrics], None
             ]:
                 # Iterate over the graph stream. For message events, yield the content.
                 # For update events, accumulate the usage metrics.

datarobot_genai/nat/agent.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import asyncio
 import logging
+from collections.abc import AsyncGenerator
 from typing import Any
 
 from nat.builder.context import Context
@@ -32,6 +33,7 @@ from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
 from datarobot_genai.core.agents.base import UsageMetrics
 from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
 
 logger = logging.getLogger(__name__)
 
@@ -160,6 +162,49 @@ class NatAgent(BaseAgent[None]):
         # Print commands may need flush=True to ensure they are displayed in real-time.
         print("Running agent with user prompt:", chat_request.messages[0].content, flush=True)
 
+        if is_streaming(completion_create_params):
+
+            async def stream_generator() -> AsyncGenerator[
+                tuple[str, MultiTurnSample | None, UsageMetrics], None
+            ]:
+                usage_metrics: UsageMetrics = {
+                    "completion_tokens": 0,
+                    "prompt_tokens": 0,
+                    "total_tokens": 0,
+                }
+                async with load_workflow(self.workflow_path) as workflow:
+                    async with workflow.run(chat_request) as runner:
+                        intermediate_future = pull_intermediate_structured()
+                        async for result in runner.result_stream():
+                            if isinstance(result, ChatResponse):
+                                result_text = result.choices[0].message.content
+                            else:
+                                result_text = str(result)
+
+                            yield (
+                                result_text,
+                                None,
+                                usage_metrics,
+                            )
+
+                        steps = await intermediate_future
+                        llm_end_steps = [
+                            step
+                            for step in steps
+                            if step.event_type == IntermediateStepType.LLM_END
+                        ]
+                        for step in llm_end_steps:
+                            if step.usage_info:
+                                token_usage = step.usage_info.token_usage
+                                usage_metrics["total_tokens"] += token_usage.total_tokens
+                                usage_metrics["prompt_tokens"] += token_usage.prompt_tokens
+                                usage_metrics["completion_tokens"] += token_usage.completion_tokens
+
+                        pipeline_interactions = create_pipeline_interactions_from_steps(steps)
+                        yield "", pipeline_interactions, usage_metrics
+
+            return stream_generator()
+
         # Create and invoke the NAT (Nemo Agent Toolkit) Agentic Workflow with the inputs
         result, steps = await self.run_nat_workflow(self.workflow_path, chat_request)
 

{datarobot_genai-0.1.62.dist-info → datarobot_genai-0.1.64.dist-info}/METADATA

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

{datarobot_genai-0.1.62.dist-info → datarobot_genai-0.1.64.dist-info}/RECORD

@@ -45,8 +45,8 @@ datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6Ke
 datarobot_genai/drmcp/core/utils.py,sha256=dSjrayWVcnC5GxQcvOIOSHaoEymPIVtG_s2ZBMlmSOw,4336
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=vCMYxwYjNDadRxSlRk6p8pHK6h_2K-PkbBTTW_lqBJ0,3318
-datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=QEySFLKdTxsXk5VoGsHobbZ_olTB9nzPGo7M99TL5mc,3894
-datarobot_genai/drmcp/core/dynamic_prompts/register.py,sha256=Q2c9zuOuywS63AcGpD4774OhQKSGHX2tpLXqW_C4JKM,6058
+datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=IEdD2Gqm4SfUdiXJB99RiWxkN6frGaxJ2SfATetMM3c,4243
+datarobot_genai/drmcp/core/dynamic_prompts/register.py,sha256=5AEh1m8GX-gPZHUdiE1VATt7IKJQk-eThcxh01sWn0I,7204
 datarobot_genai/drmcp/core/dynamic_prompts/utils.py,sha256=BZ3792AgfvYlwL0_J0MzQfGecyEA5_OKUMynEZYzCds,1136
 datarobot_genai/drmcp/core/dynamic_tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/core/dynamic_tools/register.py,sha256=3M5-F0mhUYTZJWmFDmqzsj3QAd7ut7b0kPv-JZyaTzg,9204
@@ -83,19 +83,19 @@ datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=t7f28y_ealZoA6
 datarobot_genai/drmcp/tools/predictive/project.py,sha256=KaMDAvJY4s12j_4ybA7-KcCS1yMOj-KPIKNBgCSE2iM,2536
 datarobot_genai/drmcp/tools/predictive/training.py,sha256=kxeDVLqUh9ajDk8wK7CZRRydDK8UNuTVZCB3huUihF8,23660
 datarobot_genai/langgraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/langgraph/agent.py,sha256=SpVjiA9ARqkivVVnuk3z6Vu9boE4OJqgMzIrcNQWccM,9669
+datarobot_genai/langgraph/agent.py,sha256=us-DXSgmI2271tsbzOOtDbRHdStCsOro-YRQ30ePpkA,9739
 datarobot_genai/langgraph/mcp.py,sha256=_ggxlRFkydhJgJVcAdwDWZL0WWCejctNC5X9ycAWeUM,2610
 datarobot_genai/llama_index/__init__.py,sha256=JEMkLQLuP8n14kNE3bZ2j08NdajnkJMfYjDQYqj7C0c,407
 datarobot_genai/llama_index/agent.py,sha256=V6ZsD9GcBDJS-RJo1tJtIHhyW69_78gM6_fOHFV-Piw,1829
 datarobot_genai/llama_index/base.py,sha256=MwhSlRCwsP0Z9IxaTJ4tMZO2mCMX3pvqNFl3s9aj1P8,14427
 datarobot_genai/llama_index/mcp.py,sha256=3XULVxtLBV_ItM491zWRg2KMHGszy6OTyPzHW08zO2k,2439
 datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/nat/agent.py,sha256=egp3U1KAe2nW8LpoRaMYOpNmwL6KRWNUULMoHkmKCgs,8135
+datarobot_genai/nat/agent.py,sha256=siBLDWAff2-JwZ8Q3iNpM_e4_IoSwG9IvY0hyEjNenw,10292
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=IZq_kooUL8QyDTkpEreszLQk9vCzg6-FbTjIkXR9wc0,7203
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=lOVaL_0Fl6-7GFYl3HmfqttqKpKt-2w8o92P3T7B6cU,3683
-datarobot_genai-0.1.62.dist-info/METADATA,sha256=Pyfx4lnxprmZK96X4rycMhtz-JuuvhX6b4MOJhYoYLA,5911
-datarobot_genai-0.1.62.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-datarobot_genai-0.1.62.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
-datarobot_genai-0.1.62.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.1.62.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.1.62.dist-info/RECORD,,
+datarobot_genai-0.1.64.dist-info/METADATA,sha256=xx80UKcDU-NjYz1wE6vHRUx_nJpEhX01MUi7tLU5wNs,5911
+datarobot_genai-0.1.64.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.1.64.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
+datarobot_genai-0.1.64.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.1.64.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.1.64.dist-info/RECORD,,

{datarobot_genai-0.1.62.dist-info → datarobot_genai-0.1.64.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any