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.
- {gmicloud-0.1.4 → gmicloud-0.1.5}/PKG-INFO +12 -16
- {gmicloud-0.1.4 → gmicloud-0.1.5}/README.md +11 -15
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/__init__.py +12 -1
- gmicloud-0.1.5/gmicloud/_internal/_client/_artifact_client.py +212 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_http_client.py +5 -1
- gmicloud-0.1.5/gmicloud/_internal/_client/_iam_client.py +189 -0
- gmicloud-0.1.5/gmicloud/_internal/_client/_task_client.py +153 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_enums.py +8 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/_artifact_manager.py +17 -5
- gmicloud-0.1.5/gmicloud/_internal/_manager/_iam_manager.py +36 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/_task_manager.py +19 -12
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_models.py +103 -10
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/client.py +26 -5
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/test_artifacts.py +14 -15
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/test_tasks.py +1 -1
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/PKG-INFO +12 -16
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/SOURCES.txt +1 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/pyproject.toml +1 -1
- gmicloud-0.1.4/gmicloud/_internal/_client/_artifact_client.py +0 -142
- gmicloud-0.1.4/gmicloud/_internal/_client/_iam_client.py +0 -124
- gmicloud-0.1.4/gmicloud/_internal/_client/_task_client.py +0 -108
- {gmicloud-0.1.4 → gmicloud-0.1.5}/examples/__init__.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/__init__.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/__init__.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_decorator.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_client/_file_upload_client.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_config.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_constants.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_exceptions.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/_internal/_manager/__init__.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/tests/__init__.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud/utils/uninstall_packages.py +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/dependency_links.txt +0 -0
- {gmicloud-0.1.4 → gmicloud-0.1.5}/gmicloud.egg-info/top_level.txt +0 -0
- {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.
|
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
|
-
|
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
|
-
|
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
|
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.
|
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
|
-
|
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
|
-
|
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
|
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.
|
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
|
+
}
|