datarobot-genai 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -95,11 +95,16 @@ class LLMMCPClient:
95
95
  ) -> str:
96
96
  """Call an MCP tool and return the result as a string."""
97
97
  result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
98
- return (
98
+ content = (
99
99
  result.content[0].text
100
100
  if result.content and isinstance(result.content[0], TextContent)
101
101
  else str(result.content)
102
102
  )
103
+ if result.structuredContent is not None:
104
+ structured_content = json.dumps(result.structuredContent)
105
+ else:
106
+ structured_content = ""
107
+ return f"Content: {content}\nStructured content: {structured_content}"
103
108
 
104
109
  async def _process_tool_calls(
105
110
  self,
@@ -13,9 +13,12 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import logging
16
+ from http import HTTPStatus
16
17
  from typing import Any
17
18
 
18
19
  import httpx
20
+ from pydantic import BaseModel
21
+ from pydantic import Field
19
22
 
20
23
  from .atlassian import ATLASSIAN_API_BASE
21
24
  from .atlassian import get_atlassian_cloud_id
@@ -23,10 +26,49 @@ from .atlassian import get_atlassian_cloud_id
23
26
  logger = logging.getLogger(__name__)
24
27
 
25
28
 
29
+ class _IssuePerson(BaseModel):
30
+ email_address: str = Field(alias="emailAddress")
31
+
32
+
33
+ class _IssueStatus(BaseModel):
34
+ name: str
35
+
36
+
37
+ class _IssueFields(BaseModel):
38
+ summary: str
39
+ status: _IssueStatus
40
+ reporter: _IssuePerson
41
+ assignee: _IssuePerson
42
+ created: str
43
+ updated: str
44
+
45
+
46
+ class Issue(BaseModel):
47
+ id: str
48
+ key: str
49
+ fields: _IssueFields
50
+
51
+ def as_flat_dict(self) -> dict[str, Any]:
52
+ return {
53
+ "id": self.id,
54
+ "key": self.key,
55
+ "summary": self.fields.summary,
56
+ "reporterEmailAddress": self.fields.reporter.email_address,
57
+ "assigneeEmailAddress": self.fields.assignee.email_address,
58
+ "created": self.fields.created,
59
+ "updated": self.fields.updated,
60
+ "status": self.fields.status.name,
61
+ }
62
+
63
+
26
64
  class JiraClient:
27
- """Client for interacting with Jira API using OAuth access token."""
65
+ """
66
+ Client for interacting with Jira API using OAuth access token.
28
67
 
29
- def __init__(self, access_token: str):
68
+ At the moment of creating this client, official Jira SDK is not supporting async.
69
+ """
70
+
71
+ def __init__(self, access_token: str) -> None:
30
72
  """
31
73
  Initialize Jira client with access token.
32
74
 
@@ -65,31 +107,107 @@ class JiraClient:
65
107
  self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="jira")
66
108
  return self._cloud_id
67
109
 
68
- async def get_jira_issue(self, issue_key: str) -> dict[str, Any]:
110
+ async def _get_full_url(self, url: str) -> str:
111
+ """Return URL for Jira API."""
112
+ cloud_id = await self._get_cloud_id()
113
+ return f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/{url}"
114
+
115
+ async def get_jira_issue(self, issue_key: str) -> Issue:
69
116
  """
70
117
  Get a Jira issue by its key.
71
118
 
72
119
  Args:
73
- issue_key: The key (ID) of the Jira issue, e.g., 'PROJ-123'
120
+ issue_key: The key of the Jira issue, e.g., 'PROJ-123'
74
121
 
75
122
  Returns
76
123
  -------
77
- Dictionary containing the issue data
124
+ Jira issue
78
125
 
79
126
  Raises
80
127
  ------
81
128
  httpx.HTTPStatusError: If the API request fails
82
129
  """
83
- cloud_id = await self._get_cloud_id()
84
- url = f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/issue/{issue_key}"
130
+ url = await self._get_full_url(f"issue/{issue_key}")
131
+ response = await self._client.get(
132
+ url, params={"fields": "id,key,summary,status,reporter,assignee,created,updated"}
133
+ )
134
+
135
+ if response.status_code == HTTPStatus.NOT_FOUND:
136
+ raise ValueError(f"{issue_key} not found")
137
+
138
+ response.raise_for_status()
139
+ issue = Issue(**response.json())
140
+ return issue
85
141
 
142
+ async def get_jira_issue_types(self, project_key: str) -> dict[str, str]:
143
+ """
144
+ Get Jira issue types possible for given project.
145
+
146
+ Args:
147
+ project_key: The key of the Jira project, e.g., 'PROJ'
148
+
149
+ Returns
150
+ -------
151
+ Dictionary where key is the issue type name and value is the issue type ID
152
+
153
+ Raises
154
+ ------
155
+ httpx.HTTPStatusError: If the API request fails
156
+ """
157
+ url = await self._get_full_url(f"issue/createmeta/{project_key}/issuetypes")
86
158
  response = await self._client.get(url)
159
+
87
160
  response.raise_for_status()
88
- return response.json()
161
+ jsoned = response.json()
162
+ issue_types = {
163
+ issue_type["untranslatedName"]: issue_type["id"]
164
+ for issue_type in jsoned.get("issueTypes", [])
165
+ }
166
+ return issue_types
167
+
168
+ async def create_jira_issue(
169
+ self, project_key: str, summary: str, issue_type_id: str, description: str | None
170
+ ) -> str:
171
+ """
172
+ Create Jira issue.
89
173
 
90
- async def close(self) -> None:
91
- """Close the HTTP client."""
92
- await self._client.aclose()
174
+ Args:
175
+ project_key: The key of the Jira project, e.g., 'PROJ'
176
+ summary: Summary of Jira issue (title), e.g., 'Fix bug abc'
177
+ issue_type_id: ID type of Jira issue, e.g., "1"
178
+ description: Optional description of Jira issue
179
+
180
+ Returns
181
+ -------
182
+ Jira issue key
183
+
184
+ Raises
185
+ ------
186
+ httpx.HTTPStatusError: If the API request fails
187
+ """
188
+ url = await self._get_full_url("issue")
189
+ payload = {
190
+ "fields": {
191
+ "project": {"key": project_key},
192
+ "summary": summary,
193
+ "issuetype": {"id": issue_type_id},
194
+ }
195
+ }
196
+
197
+ if description:
198
+ payload["fields"]["description"] = {
199
+ "content": [
200
+ {"content": [{"text": description, "type": "text"}], "type": "paragraph"}
201
+ ],
202
+ "type": "doc",
203
+ "version": 1,
204
+ }
205
+
206
+ response = await self._client.post(url, json=payload)
207
+
208
+ response.raise_for_status()
209
+ jsoned = response.json()
210
+ return jsoned["key"]
93
211
 
94
212
  async def __aenter__(self) -> "JiraClient":
95
213
  """Async context manager entry."""
@@ -99,4 +217,4 @@ class JiraClient:
99
217
  self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
100
218
  ) -> None:
101
219
  """Async context manager exit."""
102
- await self.close()
220
+ await self._client.aclose()
@@ -35,24 +35,67 @@ async def jira_get_issue(
35
35
 
36
36
  access_token = await get_atlassian_access_token()
37
37
  if isinstance(access_token, ToolError):
38
- return access_token
38
+ raise access_token
39
39
 
40
40
  try:
41
- client = JiraClient(
42
- access_token,
43
- )
44
- issue = await client.get_jira_issue(issue_key)
41
+ async with JiraClient(access_token) as client:
42
+ issue = await client.get_jira_issue(issue_key)
45
43
  except Exception as e:
46
- logger.error(f"Unexpected error getting Jira issue: {e}")
47
- return ToolError(
44
+ logger.error(f"Unexpected error while getting Jira issue: {e}")
45
+ raise ToolError(
48
46
  f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
49
47
  )
50
48
 
51
49
  return ToolResult(
52
50
  content=f"Successfully retrieved details for issue '{issue_key}'.",
53
- # TODO: Add more fields to the structured content, note fields here are just examples
54
- structured_content={
55
- "key": issue.get("key", issue_key),
56
- "status": issue.get("fields", {}).get("status", {}).get("name", "Unknown"),
57
- },
51
+ structured_content=issue.as_flat_dict(),
52
+ )
53
+
54
+
55
+ @dr_mcp_tool(tags={"jira", "create", "add", "issue"})
56
+ async def jira_create_issue(
57
+ *,
58
+ project_key: Annotated[str, "The key of the project where the issue should be created."],
59
+ summary: Annotated[str, "A brief summary or title for the new issue."],
60
+ issue_type: Annotated[str, "The type of issue to create (e.g., 'Task', 'Bug', 'Story')."],
61
+ description: Annotated[str | None, "Detailed description of the issue."] = None,
62
+ ) -> ToolResult:
63
+ """Create a new Jira issue with mandatory project, summary, and type information."""
64
+ if not all([project_key, summary, issue_type]):
65
+ raise ToolError(
66
+ "Argument validation error: project_key, summary, and issue_type are required fields."
67
+ )
68
+
69
+ access_token = await get_atlassian_access_token()
70
+ if isinstance(access_token, ToolError):
71
+ raise access_token
72
+
73
+ async with JiraClient(access_token) as client:
74
+ # Maybe we should cache it somehow?
75
+ # It'll be probably constant through whole mcp server lifecycle...
76
+ issue_types = await client.get_jira_issue_types(project_key=project_key)
77
+
78
+ try:
79
+ issue_type_id = issue_types[issue_type]
80
+ except KeyError:
81
+ possible_issue_types = ",".join(issue_types)
82
+ raise ToolError(
83
+ f"Unexpected issue type `{issue_type}`. Possible values are {possible_issue_types}."
84
+ )
85
+
86
+ try:
87
+ async with JiraClient(access_token) as client:
88
+ issue_key = await client.create_jira_issue(
89
+ project_key=project_key,
90
+ summary=summary,
91
+ issue_type_id=issue_type_id,
92
+ description=description,
93
+ )
94
+ except Exception as e:
95
+ logger.error(f"Unexpected error while creating Jira issue: {e}")
96
+ raise ToolError(f"An unexpected error occurred while creating Jira issue: {str(e)}")
97
+
98
+ return ToolResult(
99
+ content=f"Successfully created issue '{issue_key}'.",
100
+ structured_content={"newIssueKey": issue_key, "projectKey": project_key},
58
101
  )
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: datarobot-genai
3
- Version: 0.2.10
3
+ Version: 0.2.12
4
4
  Summary: Generic helpers for GenAI
5
5
  Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
6
6
  Author: DataRobot, Inc.
@@ -25,6 +25,8 @@ Requires-Dist: pypdf<7.0.0,>=6.1.3
25
25
  Requires-Dist: ragas<0.4.0,>=0.3.8
26
26
  Requires-Dist: requests<3.0.0,>=2.32.4
27
27
  Provides-Extra: crewai
28
+ Requires-Dist: anthropic<1.0.0,~=0.71.0; extra == 'crewai'
29
+ Requires-Dist: azure-ai-inference<2.0.0,>=1.0.0b9; extra == 'crewai'
28
30
  Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
29
31
  Requires-Dist: crewai>=1.1.0; extra == 'crewai'
30
32
  Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'
@@ -70,19 +70,19 @@ datarobot_genai/drmcp/test_utils/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT
70
70
  datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=MdoR7r3m9uT7crodyhY69yhkrM7Thpe__BBD9lB_2oA,3328
71
71
  datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rgZkPF26YCHX2FGppWE4v22l_NQ3kLSPSUimO0tD4nM,4402
72
72
  datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=0sU29Khal0CelnHBDInyTRiuPKrFFbTbIomOoUbyMhs,3271
73
- datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=Va3_5c2ToZyfIsEjK2ef5d3z-FA5SE51voikvjKPt8Q,8837
73
+ datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=TvTkDBcHscLDmqge9NhHxVo1ABtb0n4NmmG2318mQHU,9088
74
74
  datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
75
75
  datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
76
76
  datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
77
77
  datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
78
78
  datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
79
79
  datarobot_genai/drmcp/tools/clients/confluence.py,sha256=gbVxeBe7RDEEQt5UMGGW6GoAXsYLhL009dOejYIaIiQ,6325
80
- datarobot_genai/drmcp/tools/clients/jira.py,sha256=JjvssdMAWgZ3HWZkQg0a3HjpE7yz7jfRtzO4LOp47Uw,3080
80
+ datarobot_genai/drmcp/tools/clients/jira.py,sha256=aSDmw07SqpoE5fMQchb_y3Ggn4WcTUZU_1M8TwvZ3-E,6498
81
81
  datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
82
82
  datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
83
83
  datarobot_genai/drmcp/tools/confluence/tools.py,sha256=t5OqXIhUm6y9bAWymyqwEMElwTxGw1xRnkW2MgJrNF8,3106
84
84
  datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
85
- datarobot_genai/drmcp/tools/jira/tools.py,sha256=EBf8T_VwKXNticorRN8ut9ktEMS0XFeb0Uj_6TKwMio,2191
85
+ datarobot_genai/drmcp/tools/jira/tools.py,sha256=LBJkK9yjgRNZJHaqgJ3bknNnvLKpr2RLLtQYAs-O-oA,4034
86
86
  datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
87
87
  datarobot_genai/drmcp/tools/predictive/data.py,sha256=k4EJxJrl8DYVGVfJ0DM4YTfnZlC_K3OUHZ0eRUzfluI,3165
88
88
  datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
@@ -105,9 +105,9 @@ datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZ
105
105
  datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
106
106
  datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
107
107
  datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
108
- datarobot_genai-0.2.10.dist-info/METADATA,sha256=AYl5uILzc_myorpYgMkAhbm6BZIGQaaHuMNQik6rdCE,6173
109
- datarobot_genai-0.2.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
110
- datarobot_genai-0.2.10.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
111
- datarobot_genai-0.2.10.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
112
- datarobot_genai-0.2.10.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
113
- datarobot_genai-0.2.10.dist-info/RECORD,,
108
+ datarobot_genai-0.2.12.dist-info/METADATA,sha256=5DnB86Cp4uSS6x5ZjJqblA40CusRaa9V5Jw0kiGSVig,6301
109
+ datarobot_genai-0.2.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
110
+ datarobot_genai-0.2.12.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
111
+ datarobot_genai-0.2.12.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
112
+ datarobot_genai-0.2.12.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
113
+ datarobot_genai-0.2.12.dist-info/RECORD,,