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

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -23,6 +23,7 @@ from fastmcp.tools.tool import ToolResult
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 from datarobot_genai.drmcp.tools.clients.atlassian import get_atlassian_access_token
 from datarobot_genai.drmcp.tools.clients.confluence import ConfluenceClient
+from datarobot_genai.drmcp.tools.clients.confluence import ConfluenceError
 
 logger = logging.getLogger(__name__)
 
@@ -65,8 +66,8 @@ async def confluence_get_page(
                         "'space_key' is required when identifying a page by title."
                     )
                 page_response = await client.get_page_by_title(page_id_or_title, space_key)
-    except ValueError as e:
-        logger.error(f"Value error getting Confluence page: {e}")
+    except ConfluenceError as e:
+        logger.error(f"Confluence error getting page: {e}")
         raise ToolError(str(e))
     except Exception as e:
         logger.error(f"Unexpected error getting Confluence page: {e}")
@@ -79,3 +80,109 @@ async def confluence_get_page(
         content=f"Successfully retrieved page '{page_response.title}'.",
         structured_content=page_response.as_flat_dict(),
     )
+
+
+@dr_mcp_tool(tags={"confluence", "write", "create", "page"})
+async def confluence_create_page(
+    *,
+    space_key: Annotated[str, "The key of the Confluence space where the new page should live."],
+    title: Annotated[str, "The title of the new page."],
+    body_content: Annotated[
+        str,
+        "The content of the page, typically in Confluence Storage Format (XML) or raw text.",
+    ],
+    parent_id: Annotated[
+        int | None,
+        "The ID of the parent page, used to create a child page.",
+    ] = None,
+) -> ToolResult:
+    """Create a new documentation page in a specified Confluence space.
+
+    Use this tool to create new Confluence pages with content in storage format.
+    The page will be created at the root level of the space unless a parent_id
+    is provided, in which case it will be created as a child page.
+
+    Usage:
+        - Root page: space_key="PROJ", title="New Page", body_content="<p>Content</p>"
+        - Child page: space_key="PROJ", title="Sub Page", body_content="<p>Content</p>",
+                      parent_id=123456
+    """
+    if not all([space_key, title, body_content]):
+        raise ToolError(
+            "Argument validation error: space_key, title, and body_content are required fields."
+        )
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with ConfluenceClient(access_token) as client:
+            page_response = await client.create_page(
+                space_key=space_key,
+                title=title,
+                body_content=body_content,
+                parent_id=parent_id,
+            )
+    except ConfluenceError as e:
+        logger.error(f"Confluence error creating page: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error creating Confluence page: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while creating Confluence page "
+            f"'{title}' in space '{space_key}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"New page '{title}' created successfully in space '{space_key}'.",
+        structured_content={"new_page_id": page_response.page_id, "title": page_response.title},
+    )
+
+
+@dr_mcp_tool(tags={"confluence", "write", "add", "comment"})
+async def confluence_add_comment(
+    *,
+    page_id: Annotated[str, "The numeric ID of the page where the comment will be added."],
+    comment_body: Annotated[str, "The text content of the comment."],
+) -> ToolResult:
+    """Add a new comment to a specified Confluence page for collaboration.
+
+    Use this tool to add comments to Confluence pages to facilitate collaboration
+    and discussion. Comments are added at the page level.
+
+    Usage:
+        - Add comment: page_id="856391684", comment_body="Great work on this documentation!"
+    """
+    if not page_id:
+        raise ToolError("Argument validation error: 'page_id' cannot be empty.")
+
+    if not comment_body:
+        raise ToolError("Argument validation error: 'comment_body' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with ConfluenceClient(access_token) as client:
+            comment_response = await client.add_comment(
+                page_id=page_id,
+                comment_body=comment_body,
+            )
+    except ConfluenceError as e:
+        logger.error(f"Confluence error adding comment: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error adding comment to Confluence page: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while adding comment to page '{page_id}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Comment added successfully to page ID {page_id}.",
+        structured_content={
+            "comment_id": comment_response.comment_id,
+            "page_id": page_id,
+        },
+    )

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -0,0 +1,88 @@
+# 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.
+
+"""Google Drive MCP tools for interacting with Google Drive API."""
+
+import logging
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.gdrive import LIMIT
+from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
+from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveClient
+from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveError
+from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"google", "gdrive", "list", "files"})
+async def google_drive_list_files(
+    *,
+    page_size: Annotated[
+        int, f"Maximum number of files to return per page (max {MAX_PAGE_SIZE})."
+    ] = 10,
+    limit: Annotated[int, f"Total maximum number of files to return (max {LIMIT})."] = 50,
+    page_token: Annotated[
+        str | None, "The token for the next page of results, retrieved from a previous call."
+    ] = None,
+    query: Annotated[
+        str | None, "Optional filter to narrow results (e.g., 'trashed = false')."
+    ] = None,
+) -> ToolResult | ToolError:
+    """
+    List files in the user's Google Drive with pagination and filtering support.
+    Use this tool to discover file names and IDs for use with other tools.
+
+    Limit must be bigger than or equal to page size and it must be multiplication of page size.
+    Ex.
+        page size = 10 limit = 50
+        page size = 3 limit = 3
+        page size = 12 limit = 36
+    """
+    access_token = await get_gdrive_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with GoogleDriveClient(access_token) as client:
+            data = await client.list_files(
+                page_size=page_size, page_token=page_token, query=query, limit=limit
+            )
+    except GoogleDriveError as e:
+        logger.error(f"Google Drive error listing files: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error listing Google Drive files: {e}")
+        raise ToolError(f"An unexpected error occurred while listing Google Drive files: {str(e)}")
+
+    number_of_files = len(data.files)
+    next_page_info = (
+        f"Next page token needed to fetch more data: {data.next_page_token}"
+        if data.next_page_token
+        else "There're no more pages."
+    )
+    return ToolResult(
+        content=f"Successfully listed {number_of_files} files. {next_page_info}",
+        structured_content={
+            "files": [
+                file.model_dump(by_alias=True, include={"id", "name"}) for file in data.files
+            ],
+            "count": number_of_files,
+            "nextPageToken": data.next_page_token,
+        },
+    )

datarobot_genai/drmcp/tools/jira/tools.py

@@ -14,6 +14,7 @@
 
 import logging
 from typing import Annotated
+from typing import Any
 
 from fastmcp.exceptions import ToolError
 from fastmcp.tools.tool import ToolResult
@@ -25,6 +26,40 @@ from datarobot_genai.drmcp.tools.clients.jira import JiraClient
 logger = logging.getLogger(__name__)
 
 
+@dr_mcp_tool(tags={"jira", "search", "issues"})
+async def jira_search_issues(
+    *,
+    jql_query: Annotated[
+        str, "The JQL (Jira Query Language) string used to filter and search for issues."
+    ],
+    max_results: Annotated[int, "Maximum number of issues to return. Default is 50."] = 50,
+) -> ToolResult:
+    """
+    Search for Jira issues using a powerful JQL query string.
+
+    Refer to JQL documentation for advanced query construction:
+    JQL functions: https://support.atlassian.com/jira-service-management-cloud/docs/jql-functions/
+    JQL fields: https://support.atlassian.com/jira-service-management-cloud/docs/jql-fields/
+    JQL keywords: https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/
+    JQL operators: https://support.atlassian.com/jira-service-management-cloud/docs/jql-operators/
+    """
+    if not jql_query:
+        raise ToolError("Argument validation error: 'jql_query' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        issues = await client.search_jira_issues(jql_query=jql_query, max_results=max_results)
+
+    n = len(issues)
+    return ToolResult(
+        content=f"Successfully executed JQL query and retrieved {n} issue(s).",
+        structured_content={"data": [issue.as_flat_dict() for issue in issues], "count": n},
+    )
+
+
 @dr_mcp_tool(tags={"jira", "read", "get", "issue"})
 async def jira_get_issue(
     *, issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."]
@@ -99,3 +134,110 @@ async def jira_create_issue(
         content=f"Successfully created issue '{issue_key}'.",
         structured_content={"newIssueKey": issue_key, "projectKey": project_key},
     )
+
+
+@dr_mcp_tool(tags={"jira", "update", "edit", "issue"})
+async def jira_update_issue(
+    *,
+    issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."],
+    fields_to_update: Annotated[
+        dict[str, Any],
+        "A dictionary of field names and their new values (e.g., {'summary': 'New content'}).",
+    ],
+) -> ToolResult:
+    """
+    Modify descriptive fields or custom fields on an existing Jira issue using its key.
+    If you want to update issue status you should use `jira_transition_issue` tool instead.
+
+    Some fields needs very specific schema to allow update.
+    You should follow jira rest api guidance.
+    Good example is description field:
+        "description": {
+            "type": "text",
+            "version": 1,
+            "text": [
+                {
+                    "type": "paragraph",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "[HERE YOU PUT REAL DESCRIPTION]"
+                        }
+                    ]
+                }
+            ]
+        }
+    """
+    if not issue_key:
+        raise ToolError("Argument validation error: 'issue_key' cannot be empty.")
+    if not fields_to_update or not isinstance(fields_to_update, dict):
+        raise ToolError(
+            "Argument validation error: 'fields_to_update' must be a non-empty dictionary."
+        )
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with JiraClient(access_token) as client:
+            updated_fields = await client.update_jira_issue(
+                issue_key=issue_key, fields=fields_to_update
+            )
+    except Exception as e:
+        logger.error(f"Unexpected error while updating Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while updating Jira issue: {str(e)}")
+
+    updated_fields_str = ",".join(updated_fields)
+    return ToolResult(
+        content=f"Successfully updated issue '{issue_key}'. Fields modified: {updated_fields_str}.",
+        structured_content={"updatedIssueKey": issue_key, "fields": updated_fields},
+    )
+
+
+@dr_mcp_tool(tags={"jira", "update", "transition", "issue"})
+async def jira_transition_issue(
+    *,
+    issue_key: Annotated[str, "The key (ID) of the Jira issue to transition, e.g. 'PROJ-123'."],
+    transition_name: Annotated[
+        str, "The exact name of the target status/transition (e.g., 'In Progress')."
+    ],
+) -> ToolResult:
+    """
+    Move a Jira issue through its defined workflow to a new status.
+    This leverages Jira's workflow engine directly.
+    """
+    if not all([issue_key, transition_name]):
+        raise ToolError("Argument validation error: issue_key and transition name/ID are required.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        available_transitions = await client.get_available_jira_transitions(issue_key=issue_key)
+
+        try:
+            transition_id = available_transitions[transition_name]
+        except KeyError:
+            available_transitions_str = ",".join(available_transitions)
+            raise ToolError(
+                f"Unexpected transition name `{transition_name}`. "
+                f"Possible values are {available_transitions_str}."
+            )
+
+    try:
+        async with JiraClient(access_token) as client:
+            await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
+    except Exception as e:
+        logger.error(f"Unexpected error while transitioning Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while transitioning Jira issue: {str(e)}")
+
+    return ToolResult(
+        content=f"Successfully transitioned issue '{issue_key}' to status '{transition_name}'.",
+        structured_content={
+            "transitionedIssueKey": issue_key,
+            "newStatusName": transition_name,
+            "newStatusId": transition_id,
+        },
+    )

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.20.dist-info}/METADATA

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