datarobot-genai 0.2.13__py3-none-any.whl → 0.2.19__py3-none-any.whl

datarobot_genai/drmcp/tools/predictive/data.py

@@ -14,51 +14,79 @@
 
 import logging
 import os
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.core.utils import is_valid_url
 
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"data", "management", "upload"})
-async def upload_dataset_to_ai_catalog(file_path: str) -> str:
-    """
-    Upload a dataset to the DataRobot AI Catalog.
-
-    Args:
-        file_path: Path to the file to upload.
+@dr_mcp_tool(tags={"predictive", "data", "write", "upload", "catalog"})
+async def upload_dataset_to_ai_catalog(
+    file_path: Annotated[str, "The path to the dataset file to upload."] | None = None,
+    file_url: Annotated[str, "The URL to the dataset file to upload."] | None = None,
+) -> ToolError | ToolResult:
+    """Upload a dataset to the DataRobot AI Catalog / Data Registry."""
+    if not file_path and not file_url:
+        return ToolError("Either file_path or file_url must be provided.")
+    if file_path and file_url:
+        return ToolError("Please provide either file_path or file_url, not both.")
 
-    Returns
-    -------
-        A string summary of the upload result.
-    """
+    # Get client
     client = get_sdk_client()
-    if not os.path.exists(file_path):
-        logger.error(f"File not found: {file_path}")
-        return f"File not found: {file_path}"
-    catalog_item = client.Dataset.create_from_file(file_path)
-    logger.info(f"Successfully uploaded dataset: {catalog_item.id}")
-    return f"AI Catalog ID: {catalog_item.id}"
-
-
-@dr_mcp_tool(tags={"data", "management", "list"})
-async def list_ai_catalog_items() -> str:
-    """
-    List all AI Catalog items (datasets) for the authenticated user.
-
-    Returns
-    -------
-        A string summary of the AI Catalog items with their IDs and names.
-    """
+    catalog_item = None
+    # If file path is provided, create dataset from file.
+    if file_path:
+        # Does file exist?
+        if not os.path.exists(file_path):
+            logger.error("File not found: %s", file_path)
+            return ToolError(f"File not found: {file_path}")
+        catalog_item = client.Dataset.create_from_file(file_path)
+    else:
+        # Does URL exist?
+        if file_url is None or not is_valid_url(file_url):
+            logger.error("Invalid file URL: %s", file_url)
+            return ToolError(f"Invalid file URL: {file_url}")
+        catalog_item = client.Dataset.create_from_url(file_url)
+
+    if not catalog_item:
+        return ToolError("Failed to upload dataset.")
+
+    return ToolResult(
+        content=f"Successfully uploaded dataset: {catalog_item.id}",
+        structured_content={
+            "dataset_id": catalog_item.id,
+            "dataset_version_id": catalog_item.version_id,
+            "dataset_name": catalog_item.name,
+        },
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "data", "read", "list", "catalog"})
+async def list_ai_catalog_items() -> ToolResult:
+    """List all AI Catalog items (datasets) for the authenticated user."""
     client = get_sdk_client()
     datasets = client.Dataset.list()
+
     if not datasets:
         logger.info("No AI Catalog items found")
-        return "No AI Catalog items found."
-    result = "\n".join(f"{ds.id}: {ds.name}" for ds in datasets)
-    logger.info(f"Found {len(datasets)} AI Catalog items")
-    return result
+        return ToolResult(
+            content="No AI Catalog items found.",
+            structured_content={"datasets": []},
+        )
+
+    return ToolResult(
+        content=f"Found {len(datasets)} AI Catalog items.",
+        structured_content={
+            "datasets": [{"id": ds.id, "name": ds.name} for ds in datasets],
+            "count": len(datasets),
+        },
+    )
 
 
 # from fastmcp import Context

datarobot_genai/nat/agent.py

@@ -21,7 +21,6 @@ from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
 from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepType
-from nat.runtime.loader import load_workflow
 from nat.utils.type_utils import StrPath
 from openai.types.chat import CompletionCreateParams
 from ragas import MultiTurnSample
@@ -34,6 +33,8 @@ 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.core.mcp.common import MCPConfig
+from datarobot_genai.nat.helpers import load_workflow
 
 logger = logging.getLogger(__name__)
 
@@ -166,17 +167,24 @@ 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)
 
+        mcp_config = MCPConfig(
+            authorization_context=self.authorization_context,
+            forwarded_headers=self.forwarded_headers,
+        )
+        server_config = mcp_config.server_config
+        headers = server_config["headers"] if server_config else None
+
         if is_streaming(completion_create_params):
 
             async def stream_generator() -> AsyncGenerator[
                 tuple[str, MultiTurnSample | None, UsageMetrics], None
             ]:
-                usage_metrics: UsageMetrics = {
+                default_usage_metrics: UsageMetrics = {
                     "completion_tokens": 0,
                     "prompt_tokens": 0,
                     "total_tokens": 0,
                 }
-                async with load_workflow(self.workflow_path) as workflow:
+                async with load_workflow(self.workflow_path, headers=headers) as workflow:
                     async with workflow.run(chat_request) as runner:
                         intermediate_future = pull_intermediate_structured()
                         async for result in runner.result_stream():
@@ -188,7 +196,7 @@ class NatAgent(BaseAgent[None]):
                             yield (
                                 result_text,
                                 None,
-                                usage_metrics,
+                                default_usage_metrics,
                             )
 
                         steps = await intermediate_future
@@ -197,6 +205,11 @@ class NatAgent(BaseAgent[None]):
                             for step in steps
                             if step.event_type == IntermediateStepType.LLM_END
                         ]
+                        usage_metrics: UsageMetrics = {
+                            "completion_tokens": 0,
+                            "prompt_tokens": 0,
+                            "total_tokens": 0,
+                        }
                         for step in llm_end_steps:
                             if step.usage_info:
                                 token_usage = step.usage_info.token_usage
@@ -210,7 +223,7 @@ class NatAgent(BaseAgent[None]):
             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)
+        result, steps = await self.run_nat_workflow(self.workflow_path, chat_request, headers)
 
         llm_end_steps = [step for step in steps if step.event_type == IntermediateStepType.LLM_END]
         usage_metrics: UsageMetrics = {
@@ -234,7 +247,7 @@ class NatAgent(BaseAgent[None]):
         return result_text, pipeline_interactions, usage_metrics
 
     async def run_nat_workflow(
-        self, workflow_path: StrPath, chat_request: ChatRequest
+        self, workflow_path: StrPath, chat_request: ChatRequest, headers: dict[str, str] | None
     ) -> tuple[ChatResponse | str, list[IntermediateStep]]:
         """Run the NAT workflow with the provided config file and input string.
 
@@ -247,7 +260,7 @@ class NatAgent(BaseAgent[None]):
             ChatResponse | str: The result from the NAT workflow
             list[IntermediateStep]: The list of intermediate steps
         """
-        async with load_workflow(workflow_path) as workflow:
+        async with load_workflow(workflow_path, headers=headers) as workflow:
             async with workflow.run(chat_request) as runner:
                 intermediate_future = pull_intermediate_structured()
                 runner_outputs = await runner.result()

datarobot_genai/nat/helpers.py

@@ -0,0 +1,87 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# 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.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from nat.builder.workflow import Workflow
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.config import Config
+from nat.runtime.loader import PluginTypes
+from nat.runtime.loader import discover_and_register_plugins
+from nat.runtime.session import SessionManager
+from nat.utils.data_models.schema_validator import validate_schema
+from nat.utils.io.yaml_tools import yaml_load
+from nat.utils.type_utils import StrPath
+
+
+def load_config(config_file: StrPath, headers: dict[str, str] | None = None) -> Config:
+    """
+    Load a NAT configuration file with injected headers. It ensures that all plugins are
+    loaded and then validates the configuration file against the Config schema.
+
+    Parameters
+    ----------
+    config_file : StrPath
+        The path to the configuration file
+
+    Returns
+    -------
+    Config
+        The validated Config object
+    """
+    # Ensure all of the plugins are loaded
+    discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+
+    config_yaml = yaml_load(config_file)
+
+    add_headers_to_datarobot_mcp_auth(config_yaml, headers)
+
+    # Validate configuration adheres to NAT schemas
+    validated_nat_config = validate_schema(config_yaml, Config)
+
+    return validated_nat_config
+
+
+def add_headers_to_datarobot_mcp_auth(config_yaml: dict, headers: dict[str, str] | None) -> None:
+    if headers:
+        if authentication := config_yaml.get("authentication"):
+            for auth_name in authentication:
+                auth_config = authentication[auth_name]
+                if auth_config.get("_type") == "datarobot_mcp_auth":
+                    auth_config["headers"] = headers
+
+
+@asynccontextmanager
+async def load_workflow(
+    config_file: StrPath, max_concurrency: int = -1, headers: dict[str, str] | None = None
+) -> AsyncGenerator[Workflow, None]:
+    """
+    Load the NAT configuration file and create a Runner object. This is the primary entry point for
+    running NAT workflows with injected headers.
+
+    Parameters
+    ----------
+    config_file : StrPath
+        The path to the configuration file
+    max_concurrency : int, optional
+        The maximum number of parallel workflow invocations to support. Specifying 0 or -1 will
+        allow an unlimited count, by default -1
+    """
+    # Load the config object
+    config = load_config(config_file, headers=headers)
+
+    # Must yield the workflow function otherwise it cleans up
+    async with WorkflowBuilder.from_config(config=config) as workflow:
+        yield SessionManager(await workflow.build(), max_concurrency=max_concurrency)

{datarobot_genai-0.2.13.dist-info → datarobot_genai-0.2.19.dist-info}/METADATA

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

{datarobot_genai-0.2.13.dist-info → datarobot_genai-0.2.19.dist-info}/RECORD

@@ -43,7 +43,7 @@ datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_
 datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
 datarobot_genai/drmcp/core/tool_config.py,sha256=5OkC3e-vekZtdqg-DbwcyadGSrDxmZqSDey2YyGVn1M,2978
 datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
-datarobot_genai/drmcp/core/utils.py,sha256=dSjrayWVcnC5GxQcvOIOSHaoEymPIVtG_s2ZBMlmSOw,4336
+datarobot_genai/drmcp/core/utils.py,sha256=EvfpqKZ3tECMoxpIQ_tA_3rOgy6KJEYKC0lWZo_Daag,4517
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=AGJlKqgHRO0Kd7Gl-Ulw9KYBgzjTTFXWBvOUF-SuKUI,5454
 datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=4j33AKmq7kQX_EE2_RWAbP8-K5KPVEvpUginTWn_MHs,2701
@@ -67,24 +67,26 @@ datarobot_genai/drmcp/core/memory_management/__init__.py,sha256=y4yapzp3KnFMzSR6
 datarobot_genai/drmcp/core/memory_management/manager.py,sha256=gmc_SQs12YQFMWl2UbfWR40QmLV9XuCnwPZgQwKWrbA,30552
 datarobot_genai/drmcp/core/memory_management/memory_tools.py,sha256=AxzpwOlldmhhDfKZcAxaGs7Xih2SCe0XbQuXX5nQczI,6397
 datarobot_genai/drmcp/test_utils/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
-datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=MdoR7r3m9uT7crodyhY69yhkrM7Thpe__BBD9lB_2oA,3328
-datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rgZkPF26YCHX2FGppWE4v22l_NQ3kLSPSUimO0tD4nM,4402
-datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=0sU29Khal0CelnHBDInyTRiuPKrFFbTbIomOoUbyMhs,3271
-datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=TvTkDBcHscLDmqge9NhHxVo1ABtb0n4NmmG2318mQHU,9088
-datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
+datarobot_genai/drmcp/test_utils/elicitation_test_tool.py,sha256=UVKwy39nl3XcVAh6IATcN-cWL2bfrprgRQ7fbK82jTI,3287
+datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=YSk19tbaka_0ziqi7LoXie4SJs-cvi9-H00Go0ZtQWE,3575
+datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=46rH0fYYmUj7ygf968iRbdSp5u95v23BEw3Ts_c431Y,4788
+datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=sHA_BWtpgIAFp9IXiNkUeBartBMjLAauqkV9bYtCr-g,3874
+datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=YgyqHK09MB-PBaqT34heqvmvYYFtLpzzSJt7xuTJmDg,11224
+datarobot_genai/drmcp/test_utils/test_interactive.py,sha256=guXvR8q2H6VUdmvIjEJcElQJCC6lQ-oTrzbD2EkHeCs,8025
+datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=wmI-xcL0rSr56-ZoGNB8np0CZHAU583o8-Kw7fRwtMw,6232
 datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
 datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
-datarobot_genai/drmcp/tools/clients/confluence.py,sha256=gbVxeBe7RDEEQt5UMGGW6GoAXsYLhL009dOejYIaIiQ,6325
-datarobot_genai/drmcp/tools/clients/jira.py,sha256=aSDmw07SqpoE5fMQchb_y3Ggn4WcTUZU_1M8TwvZ3-E,6498
+datarobot_genai/drmcp/tools/clients/confluence.py,sha256=gDzy8t5t3b1mwEr-CuZ5BwXXQ52AXke8J_Ra7i_8T1g,13692
+datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
 datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
 datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
-datarobot_genai/drmcp/tools/confluence/tools.py,sha256=t5OqXIhUm6y9bAWymyqwEMElwTxGw1xRnkW2MgJrNF8,3106
+datarobot_genai/drmcp/tools/confluence/tools.py,sha256=jSF7yXGFqqlMcavkRIY4HbMxb7tCeunA2ST41wa2vGI,7219
 datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
-datarobot_genai/drmcp/tools/jira/tools.py,sha256=LBJkK9yjgRNZJHaqgJ3bknNnvLKpr2RLLtQYAs-O-oA,4034
+datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
-datarobot_genai/drmcp/tools/predictive/data.py,sha256=k4EJxJrl8DYVGVfJ0DM4YTfnZlC_K3OUHZ0eRUzfluI,3165
+datarobot_genai/drmcp/tools/predictive/data.py,sha256=mt3PvIel4IUAZo0HLIs6QsTtRkuzX1qg-5PHI3IJ7E8,4455
 datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
 datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
 datarobot_genai/drmcp/tools/predictive/model.py,sha256=Yih5-KedJ-1yupPLXCJsCXOdyWWi9pRvgapXDlgXWJA,4891
@@ -100,14 +102,15 @@ datarobot_genai/llama_index/agent.py,sha256=V6ZsD9GcBDJS-RJo1tJtIHhyW69_78gM6_fO
 datarobot_genai/llama_index/base.py,sha256=ovcQQtC-djD_hcLrWdn93jg23AmD6NBEj7xtw4a6K6c,14481
 datarobot_genai/llama_index/mcp.py,sha256=leXqF1C4zhuYEKFwNEfZHY4dsUuGZk3W7KArY-zxVL8,2645
 datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/nat/agent.py,sha256=jDeIS9f-8vGbeLy5gQkSjeuHINx5Fh_4BvXYERsgIIk,10516
+datarobot_genai/nat/agent.py,sha256=DuGrgqt1FzvAE-cRH_P3LTFUlwuClvbVurdwA-RsbuY,11177
 datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZC74ziNo_FHbsZgtiM,4230
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=Yu208Ed_p_4P3HdpuM7fYnKcXtimORHpKlWVPyijpU8,11356
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
 datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
-datarobot_genai-0.2.13.dist-info/METADATA,sha256=RJZ6ozRm3L6oreEu4D9gGKZLzlf3xoC7tZ3RrppBc_U,6301
-datarobot_genai-0.2.13.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.13.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.13.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.13.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.13.dist-info/RECORD,,
+datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
+datarobot_genai-0.2.19.dist-info/METADATA,sha256=rywu2LteAnHoLxfq84xckJSU10XX_sKMqli7pqsuCgg,6301
+datarobot_genai-0.2.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.19.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.19.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.19.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.19.dist-info/RECORD,,