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

datarobot_genai/drmcp/tools/clients/jira.py

@@ -107,6 +107,11 @@ class JiraClient:
         self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="jira")
         return self._cloud_id
 
+    async def _get_full_url(self, url: str) -> str:
+        """Return URL for Jira API."""
+        cloud_id = await self._get_cloud_id()
+        return f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/{url}"
+
     async def get_jira_issue(self, issue_key: str) -> Issue:
         """
         Get a Jira issue by its key.
@@ -122,9 +127,7 @@ class JiraClient:
         ------
             httpx.HTTPStatusError: If the API request fails
         """
-        cloud_id = await self._get_cloud_id()
-        url = f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/issue/{issue_key}"
-
+        url = await self._get_full_url(f"issue/{issue_key}")
         response = await self._client.get(
             url, params={"fields": "id,key,summary,status,reporter,assignee,created,updated"}
         )
@@ -136,6 +139,76 @@ class JiraClient:
         issue = Issue(**response.json())
         return issue
 
+    async def get_jira_issue_types(self, project_key: str) -> dict[str, str]:
+        """
+        Get Jira issue types possible for given project.
+
+        Args:
+            project_key: The key of the Jira project, e.g., 'PROJ'
+
+        Returns
+        -------
+            Dictionary where key is the issue type name and value is the issue type ID
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url(f"issue/createmeta/{project_key}/issuetypes")
+        response = await self._client.get(url)
+
+        response.raise_for_status()
+        jsoned = response.json()
+        issue_types = {
+            issue_type["untranslatedName"]: issue_type["id"]
+            for issue_type in jsoned.get("issueTypes", [])
+        }
+        return issue_types
+
+    async def create_jira_issue(
+        self, project_key: str, summary: str, issue_type_id: str, description: str | None
+    ) -> str:
+        """
+        Create Jira issue.
+
+        Args:
+            project_key: The key of the Jira project, e.g., 'PROJ'
+            summary: Summary of Jira issue (title), e.g., 'Fix bug abc'
+            issue_type_id: ID type of Jira issue, e.g., "1"
+            description: Optional description of Jira issue
+
+        Returns
+        -------
+            Jira issue key
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        url = await self._get_full_url("issue")
+        payload = {
+            "fields": {
+                "project": {"key": project_key},
+                "summary": summary,
+                "issuetype": {"id": issue_type_id},
+            }
+        }
+
+        if description:
+            payload["fields"]["description"] = {
+                "content": [
+                    {"content": [{"text": description, "type": "text"}], "type": "paragraph"}
+                ],
+                "type": "doc",
+                "version": 1,
+            }
+
+        response = await self._client.post(url, json=payload)
+
+        response.raise_for_status()
+        jsoned = response.json()
+        return jsoned["key"]
+
     async def __aenter__(self) -> "JiraClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/jira/tools.py

@@ -41,7 +41,7 @@ async def jira_get_issue(
         async with JiraClient(access_token) as client:
             issue = await client.get_jira_issue(issue_key)
     except Exception as e:
-        logger.error(f"Unexpected error getting Jira issue: {e}")
+        logger.error(f"Unexpected error while getting Jira issue: {e}")
         raise ToolError(
             f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
         )
@@ -50,3 +50,52 @@ async def jira_get_issue(
         content=f"Successfully retrieved details for issue '{issue_key}'.",
         structured_content=issue.as_flat_dict(),
     )
+
+
+@dr_mcp_tool(tags={"jira", "create", "add", "issue"})
+async def jira_create_issue(
+    *,
+    project_key: Annotated[str, "The key of the project where the issue should be created."],
+    summary: Annotated[str, "A brief summary or title for the new issue."],
+    issue_type: Annotated[str, "The type of issue to create (e.g., 'Task', 'Bug', 'Story')."],
+    description: Annotated[str | None, "Detailed description of the issue."] = None,
+) -> ToolResult:
+    """Create a new Jira issue with mandatory project, summary, and type information."""
+    if not all([project_key, summary, issue_type]):
+        raise ToolError(
+            "Argument validation error: project_key, summary, and issue_type are required fields."
+        )
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        # Maybe we should cache it somehow?
+        # It'll be probably constant through whole mcp server lifecycle...
+        issue_types = await client.get_jira_issue_types(project_key=project_key)
+
+        try:
+            issue_type_id = issue_types[issue_type]
+        except KeyError:
+            possible_issue_types = ",".join(issue_types)
+            raise ToolError(
+                f"Unexpected issue type `{issue_type}`. Possible values are {possible_issue_types}."
+            )
+
+    try:
+        async with JiraClient(access_token) as client:
+            issue_key = await client.create_jira_issue(
+                project_key=project_key,
+                summary=summary,
+                issue_type_id=issue_type_id,
+                description=description,
+            )
+    except Exception as e:
+        logger.error(f"Unexpected error while creating Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while creating Jira issue: {str(e)}")
+
+    return ToolResult(
+        content=f"Successfully created issue '{issue_key}'.",
+        structured_content={"newIssueKey": issue_key, "projectKey": project_key},
+    )

datarobot_genai/nat/datarobot_llm_clients.py

@@ -14,6 +14,7 @@
 
 from collections.abc import AsyncGenerator
 from typing import Any
+from typing import TypeVar
 
 from crewai import LLM
 from langchain_openai import ChatOpenAI
@@ -22,12 +23,32 @@ from llama_index.llms.litellm import LiteLLM
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+from nat.plugins.langchain.llm import (
+    _patch_llm_based_on_config as langchain_patch_llm_based_on_config,
+)
+from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 
 from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMDeploymentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMGatewayModelConfig
 from ..nat.datarobot_llm_providers import DataRobotNIMModelConfig
 
+ModelType = TypeVar("ModelType")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    if isinstance(llm_config, RetryMixin):
+        client = patch_with_retry(
+            client,
+            retries=llm_config.num_retries,
+            retry_codes=llm_config.retry_on_status_codes,
+            retry_on_messages=llm_config.retry_on_errors,
+        )
+
+    return client
+
 
 class DataRobotChatOpenAI(ChatOpenAI):
     def _get_request_payload(
@@ -77,7 +98,8 @@ async def datarobot_llm_gateway_langchain(
     config["base_url"] = config["base_url"] + "/genai/llmgw"
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    yield DataRobotChatOpenAI(**config)
+    client = DataRobotChatOpenAI(**config)
+    yield langchain_patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -90,7 +112,8 @@ async def datarobot_llm_gateway_crewai(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["base_url"] = config["base_url"].removesuffix("/api/v2")
-    yield LLM(**config)
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -103,7 +126,8 @@ async def datarobot_llm_gateway_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
-    yield DataRobotLiteLLM(**config)
+    client = DataRobotLiteLLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -119,7 +143,8 @@ async def datarobot_llm_deployment_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    yield DataRobotChatOpenAI(**config)
+    client = DataRobotChatOpenAI(**config)
+    yield langchain_patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -136,7 +161,8 @@ async def datarobot_llm_deployment_crewai(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    yield LLM(**config)
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -153,7 +179,8 @@ async def datarobot_llm_deployment_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    yield DataRobotLiteLLM(**config)
+    client = DataRobotLiteLLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -167,7 +194,8 @@ async def datarobot_nim_langchain(
     )
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
-    yield DataRobotChatOpenAI(**config)
+    client = DataRobotChatOpenAI(**config)
+    yield langchain_patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
@@ -182,7 +210,8 @@ async def datarobot_nim_crewai(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    yield LLM(**config)
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(config_type=DataRobotNIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
@@ -197,7 +226,8 @@ async def datarobot_nim_llamaindex(
     if not config["model"].startswith("datarobot/"):
         config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
-    yield DataRobotLiteLLM(**config)
+    client = DataRobotLiteLLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -212,7 +242,8 @@ async def datarobot_llm_component_langchain(
     config["stream_options"] = {"include_usage": True}
     config["model"] = config["model"].removeprefix("datarobot/")
     config.pop("use_datarobot_llm_gateway")
-    yield DataRobotChatOpenAI(**config)
+    client = DataRobotChatOpenAI(**config)
+    yield langchain_patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -229,7 +260,8 @@ async def datarobot_llm_component_crewai(
     else:
         config["api_base"] = config.pop("base_url") + "/chat/completions"
     config.pop("use_datarobot_llm_gateway")
-    yield LLM(**config)
+    client = LLM(**config)
+    yield _patch_llm_based_on_config(client, config)
 
 
 @register_llm_client(
@@ -246,4 +278,5 @@ async def datarobot_llm_component_llamaindex(
     else:
         config["api_base"] = config.pop("base_url") + "/chat/completions"
     config.pop("use_datarobot_llm_gateway")
-    yield DataRobotLiteLLM(**config)
+    client = DataRobotLiteLLM(**config)
+    yield _patch_llm_based_on_config(client, config)

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

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

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

@@ -77,12 +77,12 @@ datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosH
 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=r-sShBRTuc6nnqeFJBDu5zSU1C3mQyGnBMz9ymBSVLc,4201
+datarobot_genai/drmcp/tools/clients/jira.py,sha256=aSDmw07SqpoE5fMQchb_y3Ggn4WcTUZU_1M8TwvZ3-E,6498
 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/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
-datarobot_genai/drmcp/tools/jira/tools.py,sha256=53SYttMWriJ-9_tByE3G_ctuMvV6iHPms6La6llF9R0,1962
+datarobot_genai/drmcp/tools/jira/tools.py,sha256=LBJkK9yjgRNZJHaqgJ3bknNnvLKpr2RLLtQYAs-O-oA,4034
 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/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
@@ -102,12 +102,12 @@ datarobot_genai/llama_index/mcp.py,sha256=leXqF1C4zhuYEKFwNEfZHY4dsUuGZk3W7KArY-
 datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/nat/agent.py,sha256=jDeIS9f-8vGbeLy5gQkSjeuHINx5Fh_4BvXYERsgIIk,10516
 datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZC74ziNo_FHbsZgtiM,4230
-datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
+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.11.dist-info/METADATA,sha256=1PyKkEq0poyLVzakstjxt-wgJ4un02i0vFthu3qbv9Q,6301
-datarobot_genai-0.2.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.11.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.11.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.11.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.11.dist-info/RECORD,,
+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,,