-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathclient.py
More file actions
226 lines (193 loc) · 8.45 KB
/
Copy pathclient.py
File metadata and controls
226 lines (193 loc) · 8.45 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
import os
import asyncio
import httpx
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
class PlexeAI:
def __init__(self, api_key: Optional[str] = None, base_url: str = "https://api.plexe.ai/v0", timeout: float = 120.0):
self.api_key = api_key
if not api_key:
self.api_key = os.environ.get("PLEXE_API_KEY")
if not self.api_key:
raise ValueError("PLEXE_API_KEY must be provided or set as environment variable")
if not base_url or not isinstance(base_url, str):
raise ValueError("base_url must be a non-empty string")
if not base_url.startswith("https://"):
raise ValueError("base_url must start with 'https://' as Plexe API requires HTTPS")
self.base_url = base_url
self.client = httpx.Client(timeout=timeout)
self.async_client = httpx.AsyncClient(timeout=timeout)
def _get_headers(self) -> Dict[str, str]:
"""Get basic headers with API key."""
return {
"x-api-key": self.api_key or "",
}
def _get_json_headers(self) -> Dict[str, str]:
"""Get headers for JSON content."""
headers = self._get_headers()
headers["Content-Type"] = "application/json"
return headers
def _ensure_list(self, data_files: Union[str, Path, List[Union[str, Path]]]) -> List[Path]:
"""Convert single file path to list and ensure all paths are Path objects."""
if isinstance(data_files, (str, Path)):
data_files = [data_files]
return [Path(f) for f in data_files]
def upload_files(self, data_files: Union[str, Path, List[Union[str, Path]]]) -> str:
"""Upload data files and return upload ID."""
files = self._ensure_list(data_files)
upload_files = []
for f in files:
if not f.exists():
raise ValueError(f"File not found: {f}")
upload_files.append(('files', (f.name, open(f, 'rb'))))
response = self.client.post(
f"{self.base_url}/uploads",
files=upload_files,
headers=self._get_headers()
)
response.raise_for_status()
return response.json()["upload_id"]
async def aupload_files(self, data_files: Union[str, Path, List[Union[str, Path]]]) -> str:
"""Upload data files asynchronously."""
files = self._ensure_list(data_files)
upload_files = []
for f in files:
if not f.exists():
raise ValueError(f"File not found: {f}")
upload_files.append(('files', (f.name, open(f, 'rb'))))
response = await self.async_client.post(
f"{self.base_url}/uploads",
files=upload_files,
headers=self._get_headers()
)
response.raise_for_status()
return response.json()["upload_id"]
def build(self,
goal: str,
model_name: str,
data_files: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
upload_id: Optional[str] = None,
eval_criteria: Optional[str] = None) -> str:
"""Build a new ML model.
Args:
goal: Description of what the model should do
model_name: Name for the model
data_files: Optional path(s) to data file(s) to upload
upload_id: Optional upload_id if files were already uploaded
eval_criteria: Optional evaluation criteria
Returns:
model_version: Version ID of the created model
"""
if data_files is None and upload_id is None:
raise ValueError("Either data_files or upload_id must be provided")
if data_files is not None and upload_id is not None:
raise ValueError("Cannot provide both data_files and upload_id")
# Get upload ID - either from new upload or use provided
if data_files is not None:
upload_id = self.upload_files(data_files)
# Create model
response = self.client.post(
f"{self.base_url}/models/{model_name}/create",
json={
"upload_id": upload_id,
"goal": goal,
"eval": eval_criteria
},
headers=self._get_json_headers()
)
response.raise_for_status()
return response.json()["model_version"]
async def abuild(self,
goal: str,
model_name: str,
data_files: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
upload_id: Optional[str] = None,
eval_criteria: Optional[str] = None) -> str:
"""Async version of build()"""
if data_files is None and upload_id is None:
raise ValueError("Either data_files or upload_id must be provided")
if data_files is not None and upload_id is not None:
raise ValueError("Cannot provide both data_files and upload_id")
# Get upload ID - either from new upload or use provided
if data_files is not None:
upload_id = await self.aupload_files(data_files)
response = await self.async_client.post(
f"{self.base_url}/models/{model_name}/create",
json={
"upload_id": upload_id,
"goal": goal,
"eval": eval_criteria
},
headers=self._get_json_headers()
)
response.raise_for_status()
return response.json()["model_version"]
def get_status(self, model_name: str, model_version: str) -> Dict[str, Any]:
"""Get status of a model build."""
response = self.client.get(
f"{self.base_url}/models/{model_name}/{model_version}/status",
headers=self._get_headers()
)
response.raise_for_status()
return response.json()
async def aget_status(self, model_name: str, model_version: str) -> Dict[str, Any]:
"""Async version of get_status()"""
response = await self.async_client.get(
f"{self.base_url}/models/{model_name}/{model_version}/status",
headers=self._get_headers()
)
response.raise_for_status()
return response.json()
def infer(self, model_name: str, model_version: str, input_data: dict) -> Dict[str, Any]:
"""Run inference using a model."""
response = self.client.post(
f"{self.base_url}/models/{model_name}/{model_version}/infer",
json=input_data,
headers=self._get_json_headers()
)
response.raise_for_status()
return response.json()
async def ainfer(self, model_name: str, model_version: str, input_data: dict) -> Dict[str, Any]:
"""Async version of infer()"""
response = await self.async_client.post(
f"{self.base_url}/models/{model_name}/{model_version}/infer",
json=input_data,
headers=self._get_json_headers()
)
response.raise_for_status()
return response.json()
def batch_infer(self, model_name: str, model_version: str, inputs: List[dict]) -> List[Dict[str, Any]]:
"""Run batch predictions."""
async def run_batch():
tasks = [
self.ainfer(model_name=model_name, model_version=model_version, input_data=x)
for x in inputs
]
return await asyncio.gather(*tasks)
return asyncio.run(run_batch())
def cleanup_upload(self, upload_id: str) -> Dict[str, Any]:
"""Clean up uploaded files."""
response = self.client.delete(
f"{self.base_url}/uploads/{upload_id}",
headers=self._get_headers()
)
response.raise_for_status()
return response.json()
async def acleanup_upload(self, upload_id: str) -> Dict[str, Any]:
"""Async version of cleanup_upload()"""
response = await self.async_client.delete(
f"{self.base_url}/uploads/{upload_id}",
headers=self._get_headers()
)
response.raise_for_status()
return response.json()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.client.close()
asyncio.run(self.async_client.aclose())
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self.client.close()
await self.async_client.aclose()