gmicloud 0.1.4__tar.gz → 0.1.5__tar.gz

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.
Files changed (35) hide show
  1. {gmicloud-0.1.4 → gmicloud-0.1.5}/PKG-INFO +12 -16
  2. {gmicloud-0.1.4 → gmicloud-0.1.5}/README.md +11 -15
  3. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/__init__.py +12 -1
  4. gmicloud-0.1.5/gmicloud/_internal/_client/_artifact_client.py +212 -0
  5. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_http_client.py +5 -1
  6. gmicloud-0.1.5/gmicloud/_internal/_client/_iam_client.py +189 -0
  7. gmicloud-0.1.5/gmicloud/_internal/_client/_task_client.py +153 -0
  8. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_enums.py +8 -0
  9. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/_artifact_manager.py +17 -5
  10. gmicloud-0.1.5/gmicloud/_internal/_manager/_iam_manager.py +36 -0
  11. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/_task_manager.py +19 -12
  12. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_models.py +103 -10
  13. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/client.py +26 -5
  14. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/test_artifacts.py +14 -15
  15. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/test_tasks.py +1 -1
  16. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/PKG-INFO +12 -16
  17. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/SOURCES.txt +1 -0
  18. {gmicloud-0.1.4 → gmicloud-0.1.5}/pyproject.toml +1 -1
  19. gmicloud-0.1.4/gmicloud/_internal/_client/_artifact_client.py +0 -142
  20. gmicloud-0.1.4/gmicloud/_internal/_client/_iam_client.py +0 -124
  21. gmicloud-0.1.4/gmicloud/_internal/_client/_task_client.py +0 -108
  22. {gmicloud-0.1.4 → gmicloud-0.1.5}/examples/__init__.py +0 -0
  23. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/__init__.py +0 -0
  24. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/__init__.py +0 -0
  25. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_decorator.py +0 -0
  26. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_file_upload_client.py +0 -0
  27. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_config.py +0 -0
  28. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_constants.py +0 -0
  29. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_exceptions.py +0 -0
  30. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/__init__.py +0 -0
  31. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/__init__.py +0 -0
  32. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/utils/uninstall_packages.py +0 -0
  33. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/dependency_links.txt +0 -0
  34. {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/top_level.txt +0 -0
  35. {gmicloud-0.1.4 → gmicloud-0.1.5}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: gmicloud
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Summary: GMI Cloud Python SDK
5
5
  Author-email: GMI <gmi@gmitec.net>
6
6
  License: MIT
@@ -13,17 +13,11 @@ Description-Content-Type: text/markdown
13
13
  # GMICloud SDK (Beta)
14
14
 
15
15
  ## Overview
16
+ Before you start: Our service and GPU resource is currenly invite-only so please contact our team (getstarted@gmicloud.ai) to get invited if you don't have one yet.
16
17
 
17
- Before you start: Our service and GPU resource is currenly invite-only so please contact our team (
18
- getstarted@gmicloud.ai) to get invited if you don't have one yet.
18
+ The GMI Inference Engine SDK provides a Python interface for deploying and managing machine learning models in production environments. It allows users to create model artifacts, schedule tasks for serving models, and call inference APIs easily.
19
19
 
20
- The GMI Inference Engine SDK provides a Python interface for deploying and managing machine learning models in
21
- production environments. It allows users to create model artifacts, schedule tasks for serving models, and call
22
- inference APIs easily.
23
-
24
- This SDK streamlines the process of utilizing GMI Cloud capabilities such as deploying models with Kubernetes-based Ray
25
- services, managing resources automatically, and accessing model inference endpoints. With minimal setup, developers can
26
- focus on building ML solutions instead of infrastructure.
20
+ This SDK streamlines the process of utilizing GMI Cloud capabilities such as deploying models with Kubernetes-based Ray services, managing resources automatically, and accessing model inference endpoints. With minimal setup, developers can focus on building ML solutions instead of infrastructure.
27
21
 
28
22
  ## Features
29
23
 
@@ -51,6 +45,7 @@ Set the following environment variables:
51
45
  export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
52
46
  export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
53
47
  export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
48
+ export GMI_CLOUD_API_KEY=<YOUR_API_KEY>
54
49
  ```
55
50
 
56
51
  ### Option 2: Passing Credentials as Parameters
@@ -73,7 +68,7 @@ python -m venv venv
73
68
  source venv/bin/activate
74
69
 
75
70
  pip install -r requirements.txt
76
- python -m examples.<example_name>
71
+ python -m examples.create_task_from_artifact_template.py
77
72
  ```
78
73
 
79
74
  ### 2. Create a Task from an Artifact Template
@@ -114,11 +109,12 @@ First, you’ll retrieve all templates and create an artifact based on the desir
114
109
  ```python
115
110
  from gmicloud import *
116
111
 
112
+
117
113
  def create_artifact_from_template(client: Client) -> str:
118
114
  artifact_manager = client.artifact_manager
119
115
 
120
116
  # Get all artifact templates
121
- templates = artifact_manager.get_artifact_templates()
117
+ templates = artifact_manager.get_public_templates()
122
118
  for template in templates:
123
119
  if template.artifact_template_id == "qwen_2.5_14b_instruct_template_001":
124
120
  # Create an artifact from a template
@@ -235,10 +231,10 @@ password: Optional[str] = ""
235
231
 
236
232
  ## Notes & Troubleshooting
237
233
 
238
- Ensure Credentials are Correct: Double-check your environment variables or parameters passed into the Client object.
239
- Artifact Status: It may take a few minutes for an artifact or task to transition to the "running" state.
240
- Inference Endpoint Readiness: Use the task endpoint only after the task status changes to "running".
241
- Default OpenAI Key: By default, the OpenAI API base URL is derived from the endpoint provided by GMI.
234
+ * Ensure Credentials are Correct: Double-check your environment variables or parameters passed into the Client object.
235
+ * Artifact Status: It may take a few minutes for an artifact or task to transition to the "running" state.
236
+ * Inference Endpoint Readiness: Use the task endpoint only after the task status changes to "running".
237
+ * Default OpenAI Key: By default, the OpenAI API base URL is derived from the endpoint provided by GMI.
242
238
 
243
239
  ## Contributing
244
240
 
@@ -1,17 +1,11 @@
1
1
  # GMICloud SDK (Beta)
2
2
 
3
3
  ## Overview
4
+ Before you start: Our service and GPU resource is currenly invite-only so please contact our team (getstarted@gmicloud.ai) to get invited if you don't have one yet.
4
5
 
5
- Before you start: Our service and GPU resource is currenly invite-only so please contact our team (
6
- getstarted@gmicloud.ai) to get invited if you don't have one yet.
6
+ The GMI Inference Engine SDK provides a Python interface for deploying and managing machine learning models in production environments. It allows users to create model artifacts, schedule tasks for serving models, and call inference APIs easily.
7
7
 
8
- The GMI Inference Engine SDK provides a Python interface for deploying and managing machine learning models in
9
- production environments. It allows users to create model artifacts, schedule tasks for serving models, and call
10
- inference APIs easily.
11
-
12
- This SDK streamlines the process of utilizing GMI Cloud capabilities such as deploying models with Kubernetes-based Ray
13
- services, managing resources automatically, and accessing model inference endpoints. With minimal setup, developers can
14
- focus on building ML solutions instead of infrastructure.
8
+ This SDK streamlines the process of utilizing GMI Cloud capabilities such as deploying models with Kubernetes-based Ray services, managing resources automatically, and accessing model inference endpoints. With minimal setup, developers can focus on building ML solutions instead of infrastructure.
15
9
 
16
10
  ## Features
17
11
 
@@ -39,6 +33,7 @@ Set the following environment variables:
39
33
  export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
40
34
  export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
41
35
  export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
36
+ export GMI_CLOUD_API_KEY=<YOUR_API_KEY>
42
37
  ```
43
38
 
44
39
  ### Option 2: Passing Credentials as Parameters
@@ -61,7 +56,7 @@ python -m venv venv
61
56
  source venv/bin/activate
62
57
 
63
58
  pip install -r requirements.txt
64
- python -m examples.<example_name>
59
+ python -m examples.create_task_from_artifact_template.py
65
60
  ```
66
61
 
67
62
  ### 2. Create a Task from an Artifact Template
@@ -102,11 +97,12 @@ First, you’ll retrieve all templates and create an artifact based on the desir
102
97
  ```python
103
98
  from gmicloud import *
104
99
 
100
+
105
101
  def create_artifact_from_template(client: Client) -> str:
106
102
  artifact_manager = client.artifact_manager
107
103
 
108
104
  # Get all artifact templates
109
- templates = artifact_manager.get_artifact_templates()
105
+ templates = artifact_manager.get_public_templates()
110
106
  for template in templates:
111
107
  if template.artifact_template_id == "qwen_2.5_14b_instruct_template_001":
112
108
  # Create an artifact from a template
@@ -223,10 +219,10 @@ password: Optional[str] = ""
223
219
 
224
220
  ## Notes & Troubleshooting
225
221
 
226
- Ensure Credentials are Correct: Double-check your environment variables or parameters passed into the Client object.
227
- Artifact Status: It may take a few minutes for an artifact or task to transition to the "running" state.
228
- Inference Endpoint Readiness: Use the task endpoint only after the task status changes to "running".
229
- Default OpenAI Key: By default, the OpenAI API base URL is derived from the endpoint provided by GMI.
222
+ * Ensure Credentials are Correct: Double-check your environment variables or parameters passed into the Client object.
223
+ * Artifact Status: It may take a few minutes for an artifact or task to transition to the "running" state.
224
+ * Inference Endpoint Readiness: Use the task endpoint only after the task status changes to "running".
225
+ * Default OpenAI Key: By default, the OpenAI API base URL is derived from the endpoint provided by GMI.
230
226
 
231
227
  ## Contributing
232
228
 
@@ -1,3 +1,6 @@
1
+ import logging
2
+ import os
3
+
1
4
  from ._internal._models import (
2
5
  Artifact,
3
6
  ArtifactData,
@@ -16,7 +19,8 @@ from ._internal._models import (
16
19
  )
17
20
  from ._internal._enums import (
18
21
  BuildStatus,
19
- TaskEndpointStatus
22
+ TaskEndpointStatus,
23
+ TaskStatus
20
24
  )
21
25
  from .client import Client
22
26
 
@@ -39,3 +43,10 @@ __all__ = [
39
43
  "BuildStatus",
40
44
  "TaskEndpointStatus",
41
45
  ]
46
+
47
+ # Configure logging
48
+ log_level = os.getenv("GMI_CLOUD_LOG_LEVEL", "INFO").upper()
49
+ logging.basicConfig(
50
+ level=log_level,
51
+ format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
52
+ )
@@ -0,0 +1,212 @@
1
+ from typing import List
2
+ import logging
3
+ from requests.exceptions import RequestException
4
+
5
+ from ._http_client import HTTPClient
6
+ from ._iam_client import IAMClient
7
+ from ._decorator import handle_refresh_token
8
+ from .._models import *
9
+ from .._config import ARTIFACT_SERVICE_BASE_URL
10
+
11
+ logger = logging.getLogger(__name__)
12
+
13
+
14
+ class ArtifactClient:
15
+ """
16
+ Client for interacting with the Artifact Service API.
17
+
18
+ This client provides methods to perform CRUD operations on artifacts,
19
+ as well as generating signed URLs for uploading large files.
20
+ """
21
+
22
+ def __init__(self, iam_client: IAMClient):
23
+ """
24
+ Initializes the ArtifactClient with an HTTPClient configured
25
+ to communicate with the Artifact Service base URL.
26
+ """
27
+ self.client = HTTPClient(ARTIFACT_SERVICE_BASE_URL)
28
+ self.iam_client = iam_client
29
+
30
+ @handle_refresh_token
31
+ def get_artifact(self, artifact_id: str) -> Optional[Artifact]:
32
+ """
33
+ Fetches an artifact by its ID.
34
+
35
+ :param artifact_id: The ID of the artifact to fetch.
36
+ :return: The Artifact object or None if an error occurs.
37
+ """
38
+ try:
39
+ response = self.client.get(
40
+ "/get_artifact",
41
+ self.iam_client.get_custom_headers(),
42
+ {"artifact_id": artifact_id}
43
+ )
44
+ return Artifact.model_validate(response) if response else None
45
+ except (RequestException, ValueError) as e:
46
+ logger.error(f"Failed to fetch artifact {artifact_id}: {e}")
47
+ return None
48
+
49
+ @handle_refresh_token
50
+ def get_all_artifacts(self) -> List[Artifact]:
51
+ """
52
+ Fetches all artifacts.
53
+
54
+ :return: A list of Artifact objects. If an error occurs, returns an empty list.
55
+ """
56
+ try:
57
+ response = self.client.get("/get_all_artifacts", self.iam_client.get_custom_headers())
58
+ if not response:
59
+ logger.error("Empty response from /get_all_artifacts")
60
+ return []
61
+ return [Artifact.model_validate(item) for item in response]
62
+ except (RequestException, ValueError) as e:
63
+ logger.error(f"Failed to fetch all artifacts: {e}")
64
+ return []
65
+
66
+ @handle_refresh_token
67
+ def create_artifact(self, request: CreateArtifactRequest) -> Optional[CreateArtifactResponse]:
68
+ """
69
+ Creates a new artifact in the service.
70
+
71
+ :param request: The request object containing artifact details.
72
+ :return: The response object containing the created artifact details, or None on error.
73
+ """
74
+ try:
75
+ response = self.client.post(
76
+ "/create_artifact",
77
+ self.iam_client.get_custom_headers(),
78
+ request.model_dump()
79
+ )
80
+ return CreateArtifactResponse.model_validate(response) if response else None
81
+ except (RequestException, ValueError) as e:
82
+ logger.error(f"Failed to create artifact: {e}")
83
+ return None
84
+
85
+ @handle_refresh_token
86
+ def create_artifact_from_template(self, artifact_template_id: str) -> Optional[CreateArtifactFromTemplateResponse]:
87
+ """
88
+ Creates a new artifact in the service.
89
+
90
+ :param artifact_template_id: The ID of the artifact template to use.
91
+ :return: The response object containing the created artifact details or None if an error occurs.
92
+ """
93
+ try:
94
+ response = self.client.post(
95
+ "/create_artifact_from_template",
96
+ self.iam_client.get_custom_headers(),
97
+ {"artifact_template_id": artifact_template_id}
98
+ )
99
+ return CreateArtifactFromTemplateResponse.model_validate(response) if response else None
100
+ except (RequestException, ValueError) as e:
101
+ logger.error(f"Failed to create artifact from template {artifact_template_id}: {e}")
102
+ return None
103
+
104
+ @handle_refresh_token
105
+ def rebuild_artifact(self, artifact_id: str) -> Optional[RebuildArtifactResponse]:
106
+ """
107
+ Rebuilds an artifact in the service.
108
+
109
+ :param artifact_id: The ID of the artifact to rebuild.
110
+ :return: The response object containing the rebuilt artifact details or None if an error occurs.
111
+ """
112
+ try:
113
+ response = self.client.post(
114
+ "/rebuild_artifact",
115
+ self.iam_client.get_custom_headers(),
116
+ {"artifact_id": artifact_id}
117
+ )
118
+ return RebuildArtifactResponse.model_validate(response) if response else None
119
+ except (RequestException, ValueError) as e:
120
+ logger.error(f"Failed to rebuild artifact {artifact_id}: {e}")
121
+ return None
122
+
123
+ @handle_refresh_token
124
+ def delete_artifact(self, artifact_id: str) -> Optional[DeleteArtifactResponse]:
125
+ """
126
+ Deletes an artifact by its ID.
127
+
128
+ :param artifact_id: The ID of the artifact to delete.
129
+ :return: The response object containing the deleted artifact details or None if an error occurs.
130
+ """
131
+ try:
132
+ response = self.client.delete(
133
+ "/delete_artifact",
134
+ self.iam_client.get_custom_headers(),
135
+ {"artifact_id": artifact_id}
136
+ )
137
+ return DeleteArtifactResponse.model_validate(response) if response else None
138
+ except (RequestException, ValueError) as e:
139
+ logger.error(f"Failed to delete artifact {artifact_id}: {e}")
140
+ return None
141
+
142
+ @handle_refresh_token
143
+ def get_bigfile_upload_url(self, request: GetBigFileUploadUrlRequest) -> Optional[GetBigFileUploadUrlResponse]:
144
+ """
145
+ Generates a pre-signed URL for uploading a large file.
146
+
147
+ :param request: The request object containing the artifact ID, file name, and file type.
148
+ :return: The response object containing the pre-signed URL and upload details, or None if an error occurs.
149
+ """
150
+ try:
151
+ response = self.client.post("/get_bigfile_upload_url",
152
+ self.iam_client.get_custom_headers(),
153
+ request.model_dump())
154
+
155
+ if not response:
156
+ logger.error("Empty response from /get_bigfile_upload_url")
157
+ return None
158
+
159
+ return GetBigFileUploadUrlResponse.model_validate(response)
160
+
161
+ except (RequestException, ValueError) as e:
162
+ logger.error(f"Failed to generate upload URL: {e}")
163
+ return None
164
+
165
+ @handle_refresh_token
166
+ def delete_bigfile(self, request: DeleteBigfileRequest) -> Optional[DeleteBigfileResponse]:
167
+ """
168
+ Deletes a large file associated with an artifact.
169
+
170
+ :param request: The request object containing the artifact ID and file name.
171
+ :return: The response object containing the deletion status, or None if an error occurs.
172
+ """
173
+ try:
174
+ response = self.client.delete("/delete_bigfile",
175
+ self.iam_client.get_custom_headers(),
176
+ request.model_dump())
177
+
178
+ if not response:
179
+ logger.error("Empty response from /delete_bigfile")
180
+ return None
181
+
182
+ return DeleteBigfileResponse.model_validate(response)
183
+
184
+ except (RequestException, ValueError) as e:
185
+ logger.error(f"Failed to delete big file: {e}")
186
+ return None
187
+
188
+ @handle_refresh_token
189
+ def get_public_templates(self) -> List[ArtifactTemplate]:
190
+ """
191
+ Fetches all artifact templates.
192
+
193
+ :return: A list of ArtifactTemplate objects.
194
+ :rtype: List[ArtifactTemplate]
195
+ """
196
+ try:
197
+ response = self.client.get("/get_public_templates", self.iam_client.get_custom_headers())
198
+
199
+ if not response:
200
+ logger.error("Empty response received from /get_public_templates API")
201
+ return []
202
+
203
+ try:
204
+ result = GetPublicTemplatesResponse.model_validate(response)
205
+ return result.artifact_templates
206
+ except ValueError as ve:
207
+ logger.error(f"Failed to validate response data: {ve}")
208
+ return []
209
+
210
+ except RequestException as e:
211
+ logger.error(f"Request to /get_public_templates failed: {e}")
212
+ return []
@@ -1,8 +1,12 @@
1
+ import logging
2
+
1
3
  import requests
2
4
  from .._exceptions import APIError
3
5
  from .._exceptions import UnauthorizedError
4
6
  from .._constants import *
5
7
 
8
+ logger = logging.getLogger(__name__)
9
+
6
10
 
7
11
  class HTTPClient:
8
12
  """
@@ -51,6 +55,7 @@ class HTTPClient:
51
55
  response = None
52
56
  try:
53
57
  response = requests.request(method, url, params=params, json=data, headers=headers)
58
+ logger.debug(response.text)
54
59
  if response.status_code == 401:
55
60
  raise UnauthorizedError("Access token expired or invalid.")
56
61
  elif response.status_code != 200 and response.status_code != 201:
@@ -61,7 +66,6 @@ class HTTPClient:
61
66
  raise APIError(f"HTTP Request failed: {error_message}")
62
67
  # Raise for HTTP errors
63
68
  response.raise_for_status()
64
- print(response.text)
65
69
 
66
70
  except requests.exceptions.RequestException as e:
67
71
  raise APIError(f"HTTP Request failed: {str(e)}")
@@ -0,0 +1,189 @@
1
+ import jwt
2
+ import logging
3
+ from requests.exceptions import RequestException
4
+
5
+ from ._http_client import HTTPClient
6
+ from .._config import IAM_SERVICE_BASE_URL
7
+ from .._models import *
8
+ from .._constants import CLIENT_ID_HEADER, AUTHORIZATION_HEADER
9
+
10
+ logger = logging.getLogger(__name__)
11
+
12
+
13
+ class IAMClient:
14
+ """
15
+ Client for interacting with the IAM Service API.
16
+ """
17
+
18
+ def __init__(self, client_id: str, email: str, password: str):
19
+ """
20
+ Initialize IAMClient with client credentials and an IAMClient instance.
21
+
22
+ :param client_id: Client ID for login.
23
+ :param email: Email for login.
24
+ :param password: Password for login.
25
+ """
26
+ self._client_id = client_id
27
+ self._email = email
28
+ self._password = password
29
+ self._access_token = ""
30
+ self._refresh_token = ""
31
+ self._user_id = ""
32
+ self._organization_id = ""
33
+ self.client = HTTPClient(IAM_SERVICE_BASE_URL)
34
+
35
+ def login(self) -> bool:
36
+ """
37
+ Logs in a user with the given email and password.
38
+ Returns True if login is successful, otherwise False.
39
+ """
40
+ try:
41
+ custom_headers = {CLIENT_ID_HEADER: self._client_id}
42
+ req = AuthTokenRequest(email=self._email, password=self._password)
43
+ auth_tokens_result = self.client.post("/me/auth-tokens", custom_headers, req.model_dump())
44
+
45
+ if not auth_tokens_result:
46
+ logger.error("Login failed: Received empty response from auth-tokens endpoint")
47
+ return False
48
+
49
+ auth_tokens_resp = AuthTokenResponse.model_validate(auth_tokens_result)
50
+
51
+ # Handle 2FA
52
+ if auth_tokens_resp.is2FARequired:
53
+ for attempt in range(3):
54
+ code = input(f"Attempt {attempt + 1}/3: Please enter the 2FA code: ")
55
+ create_session_req = CreateSessionRequest(
56
+ type="native", authToken=auth_tokens_resp.authToken, otpCode=code
57
+ )
58
+ try:
59
+ session_result = self.client.post("/me/sessions", custom_headers,
60
+ create_session_req.model_dump())
61
+ if session_result:
62
+ break
63
+ except RequestException:
64
+ logger.warning("Invalid 2FA code, please try again.")
65
+ if attempt == 2:
66
+ logger.error("Failed to create session after 3 incorrect 2FA attempts.")
67
+ return False
68
+ else:
69
+ create_session_req = CreateSessionRequest(type="native", authToken=auth_tokens_resp.authToken,
70
+ otpCode=None)
71
+ session_result = self.client.post("/me/sessions", custom_headers, create_session_req.model_dump())
72
+
73
+ create_session_resp = CreateSessionResponse.model_validate(session_result)
74
+
75
+ self._access_token = create_session_resp.accessToken
76
+ self._refresh_token = create_session_resp.refreshToken
77
+ self._user_id = self.parse_user_id()
78
+
79
+ # Fetch profile to get organization ID
80
+ profile_result = self.client.get("/me/profile", self.get_custom_headers())
81
+ if not profile_result:
82
+ logger.error("Failed to fetch user profile data.")
83
+ return False
84
+
85
+ profile_resp = ProfileResponse.model_validate(profile_result)
86
+ self._organization_id = profile_resp.organization.id
87
+
88
+ return True
89
+ except (RequestException, ValueError, KeyError) as e:
90
+ logger.error(f"Login failed due to exception: {e}")
91
+ return False
92
+
93
+ def refresh_token(self) -> bool:
94
+ """
95
+ Refreshes the access token. Returns True on success, False otherwise.
96
+ """
97
+ try:
98
+ custom_headers = {CLIENT_ID_HEADER: self._client_id}
99
+ result = self.client.patch("/me/sessions", custom_headers, {"refreshToken": self._refresh_token})
100
+
101
+ if not result:
102
+ logger.error("Failed to refresh token: Empty response received")
103
+ return False
104
+
105
+ resp = CreateSessionResponse.model_validate(result)
106
+ self._access_token = resp.accessToken
107
+ self._refresh_token = resp.refreshToken
108
+
109
+ return True
110
+ except (RequestException, ValueError) as e:
111
+ logger.error(f"Token refresh failed: {e}")
112
+ return False
113
+
114
+ def create_org_api_key(self, request: CreateAPIKeyRequest) -> Optional[str]:
115
+ """
116
+ Creates a new API key for the current user.
117
+ """
118
+ try:
119
+ result = self.client.post(f"/organizations/{self.get_organization_id()}/api-keys",
120
+ self.get_custom_headers(), request.model_dump())
121
+
122
+ return CreateAPIKeyResponse.model_validate(result).key if result else None
123
+ except (RequestException, ValueError) as e:
124
+ logger.error(f"Failed to create API key: {e}")
125
+ return None
126
+
127
+ def get_org_api_keys(self) -> Optional[GetAPIKeysResponse]:
128
+ """
129
+ Fetches all API keys for the current user.
130
+ """
131
+ try:
132
+ result = self.client.get(f"/organizations/{self.get_organization_id()}/api-keys",
133
+ self.get_custom_headers())
134
+
135
+ return GetAPIKeysResponse.model_validate(result) if result else None
136
+ except (RequestException, ValueError) as e:
137
+ logger.error(f"Failed to retrieve organization API keys: {e}")
138
+ return None
139
+
140
+ def parse_user_id(self) -> str:
141
+ """
142
+ Parses the current access token and returns the user ID.
143
+ """
144
+ return self.parse_token()["userId"]
145
+
146
+ def parse_token(self) -> dict:
147
+ """
148
+ Parses the current access token and returns the payload as a dictionary.
149
+ """
150
+ return jwt.decode(self._access_token, options={"verify_signature": False})
151
+
152
+ def get_access_token(self) -> str:
153
+ """
154
+ Gets the current access token.
155
+ """
156
+ return self._access_token
157
+
158
+ def get_refresh_token(self) -> str:
159
+ """
160
+ Gets the current refresh token.
161
+ """
162
+ return self._refresh_token
163
+
164
+ def get_user_id(self) -> str:
165
+ """
166
+ Gets the current user ID.
167
+ """
168
+ return self._user_id
169
+
170
+ def get_client_id(self) -> str:
171
+ """
172
+ Gets the current client ID.
173
+ """
174
+ return self._client_id
175
+
176
+ def get_organization_id(self) -> str:
177
+ """
178
+ Gets the current organization ID.
179
+ """
180
+ return self._organization_id
181
+
182
+ def get_custom_headers(self) -> dict:
183
+ """
184
+ Gets the custom headers for the IAM client.
185
+ """
186
+ return {
187
+ AUTHORIZATION_HEADER: f'Bearer {self._access_token}',
188
+ CLIENT_ID_HEADER: self._client_id
189
+ }