335 lines
12 KiB
Python
335 lines
12 KiB
Python
"""
|
|
API Client for Asset Manager
|
|
This module handles communication with the external API
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
import os
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
class AssetManagerAPIClient:
|
|
"""Client for communicating with the Asset Manager API"""
|
|
|
|
def __init__(self, api_key: str, base_url: str = None):
|
|
"""
|
|
Initialize the API client
|
|
|
|
Args:
|
|
api_key: API authentication key
|
|
base_url: Base URL for the API (optional, can be set via config)
|
|
"""
|
|
self.api_key = api_key
|
|
self.base_url = base_url or "https://api.folders.nodehaus.io/api"
|
|
self.upload_base_url = "https://api.upload.nodehaus.io"
|
|
self.headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"x-api-key": api_key,
|
|
}
|
|
self.timeout = 30
|
|
self.chunk_size = 10 * 1024 * 1024 #10MB chunks
|
|
|
|
def get_organizations(self) -> Dict:
|
|
"""
|
|
Fetch list of organizations from API
|
|
|
|
Returns:
|
|
Dictionary containing status and organizations list
|
|
"""
|
|
try:
|
|
url = f"{self.base_url}/organizations"
|
|
|
|
response = requests.get(url, headers=self.headers, timeout=self.timeout)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
if data.get("success") and data.get("organizations"):
|
|
return {
|
|
"status": "success",
|
|
"organizations": data.get("organizations", []),
|
|
}
|
|
else:
|
|
return {
|
|
"status": "error",
|
|
"message": data.get("error", "Failed to load organizations"),
|
|
"organizations": [],
|
|
}
|
|
else:
|
|
return {
|
|
"status": "error",
|
|
"message": f"HTTP error! status: {response.status_code}",
|
|
"organizations": [],
|
|
}
|
|
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e), "organizations": []}
|
|
|
|
def get_projects(self, organization_id: str = None) -> Dict:
|
|
"""
|
|
Fetch list of projects for an organization
|
|
|
|
Args:
|
|
organization_id: Organization ID (optional)
|
|
|
|
Returns:
|
|
Dictionary containing status and projects list
|
|
"""
|
|
try:
|
|
url = f"{self.base_url}/projects"
|
|
|
|
params = {}
|
|
if organization_id:
|
|
params["organization_id"] = organization_id
|
|
|
|
response = requests.get(
|
|
url, headers=self.headers, params=params, timeout=self.timeout
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
if data.get("success") and data.get("folders"):
|
|
return {"status": "success", "projects": data.get("folders", [])}
|
|
else:
|
|
return {
|
|
"status": "error",
|
|
"message": data.get("error", "Failed to load projects"),
|
|
"projects": [],
|
|
}
|
|
else:
|
|
return {
|
|
"status": "error",
|
|
"message": f"HTTP error! status: {response.status_code}",
|
|
"projects": [],
|
|
}
|
|
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e), "projects": []}
|
|
|
|
def upload_asset(
|
|
self,
|
|
file_path: str,
|
|
project_id: str,
|
|
folder_id: str = None,
|
|
organization_id: str = None,
|
|
metadata: Optional[Dict] = None,
|
|
) -> Dict:
|
|
"""
|
|
Upload a single asset using multipart upload
|
|
|
|
Args:
|
|
file_path: Path to the file to upload
|
|
project_id: Target project ID
|
|
folder_id: Target folder ID (defaults to project_id if not provided)
|
|
organization_id: Organization ID
|
|
metadata: Optional metadata dictionary (ComfyUI workflow JSON)
|
|
|
|
Returns:
|
|
Dictionary containing upload status and result
|
|
"""
|
|
upload_id = None
|
|
key = None
|
|
|
|
try:
|
|
file_size = os.path.getsize(file_path)
|
|
filename = os.path.basename(file_path)
|
|
file_ext = os.path.splitext(file_path)[1].lower()
|
|
|
|
type_mapping = {
|
|
".png": ("image", "image/png"),
|
|
".jpg": ("image", "image/jpeg"),
|
|
".jpeg": ("image", "image/jpeg"),
|
|
".svg": ("image", "image/svg+xml"),
|
|
".gif": ("image", "image/gif"),
|
|
".webp": ("image", "image/webp"),
|
|
".bmp": ("image", "image/bmp"),
|
|
".mp4": ("video", "video/mp4"),
|
|
".mov": ("video", "video/quicktime"),
|
|
".avi": ("video", "video/x-msvideo"),
|
|
".mkv": ("video", "video/x-matroska"),
|
|
".txt": ("text", "text/plain"),
|
|
".json": ("text", "application/json"),
|
|
".mp3": ("audio", "audio/mpeg"),
|
|
".wav": ("audio", "audio/wav"),
|
|
".flac": ("audio", "audio/flac"),
|
|
".obj": ("3D", "model/obj"),
|
|
".fbx": ("3D", "model/fbx"),
|
|
".gltf": ("3D", "model/gltf+json"),
|
|
".glb": ("3D", "model/gltf-binary"),
|
|
}
|
|
asset_type, content_type = type_mapping.get(
|
|
file_ext, ("image", "application/octet-stream")
|
|
)
|
|
|
|
create_response = requests.post(
|
|
f"{self.upload_base_url}/api/upload/create",
|
|
headers={"Content-Type": "application/json", "x-api-key": self.api_key},
|
|
json={
|
|
"organizationId": organization_id,
|
|
"projectId": project_id,
|
|
"platform": "comfyui",
|
|
"fileName": filename,
|
|
"fileSize": file_size,
|
|
"contentType": content_type,
|
|
"type": asset_type,
|
|
"title": f"comfyui_{filename}",
|
|
"folderId": folder_id or project_id,
|
|
"metadata": metadata,
|
|
},
|
|
timeout=self.timeout,
|
|
)
|
|
create_response.raise_for_status()
|
|
create_result = create_response.json()
|
|
upload_id = create_result["uploadId"]
|
|
key = create_result["key"]
|
|
|
|
parts = []
|
|
with open(file_path, "rb") as f:
|
|
part_number = 1
|
|
while True:
|
|
chunk = f.read(self.chunk_size)
|
|
if not chunk:
|
|
break
|
|
|
|
part_response = requests.post(
|
|
f"{self.upload_base_url}/api/upload/part",
|
|
headers={
|
|
"x-api-key": self.api_key,
|
|
"X-Upload-Id": upload_id,
|
|
"X-Part-Number": str(part_number),
|
|
"X-Key": key,
|
|
},
|
|
data=chunk,
|
|
timeout=self.timeout * 2, #double timeout for large chunks
|
|
)
|
|
part_response.raise_for_status()
|
|
part_result = part_response.json()
|
|
parts.append(
|
|
{
|
|
"partNumber": part_result["partNumber"],
|
|
"etag": part_result["etag"],
|
|
}
|
|
)
|
|
part_number += 1
|
|
|
|
complete_response = requests.post(
|
|
f"{self.upload_base_url}/api/upload/complete",
|
|
headers={"Content-Type": "application/json", "x-api-key": self.api_key},
|
|
json={
|
|
"uploadId": upload_id,
|
|
"key": key,
|
|
"parts": parts,
|
|
"type": asset_type,
|
|
"title": f"comfyui_{filename}",
|
|
"folderId": folder_id or project_id,
|
|
"organizationId": organization_id,
|
|
"projectId": project_id,
|
|
"platform": "comfyui",
|
|
"metadata": metadata if metadata else {},
|
|
},
|
|
timeout=self.timeout,
|
|
)
|
|
complete_response.raise_for_status()
|
|
complete_result = complete_response.json()
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": f"Successfully uploaded {filename}",
|
|
"data": complete_result,
|
|
}
|
|
|
|
except Exception as e:
|
|
if upload_id and key:
|
|
try:
|
|
requests.post(
|
|
f"{self.upload_base_url}/api/upload/abort",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"x-api-key": self.api_key,
|
|
},
|
|
json={"uploadId": upload_id, "key": key},
|
|
timeout=self.timeout,
|
|
)
|
|
except Exception as abort_error:
|
|
pass
|
|
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
return {"status": "error", "message": f"Upload failed: {str(e)}"}
|
|
|
|
def batch_upload_assets(
|
|
self,
|
|
file_paths: List[str],
|
|
project_id: str,
|
|
folder_id: str = None,
|
|
organization_id: str = None,
|
|
metadata: Optional[Dict] = None,
|
|
) -> Dict:
|
|
"""
|
|
Upload multiple assets in batches
|
|
|
|
Args:
|
|
file_paths: List of file paths to upload
|
|
project_id: Target project ID
|
|
folder_id: Target folder ID (defaults to project_id if not provided)
|
|
organization_id: Organization ID
|
|
metadata: Optional metadata dictionary
|
|
|
|
Returns:
|
|
Dictionary containing batch upload status and results
|
|
"""
|
|
results = {
|
|
"status": "success",
|
|
"total": len(file_paths),
|
|
"successful": 0,
|
|
"failed": 0,
|
|
"errors": [],
|
|
}
|
|
|
|
for file_path in file_paths:
|
|
result = self.upload_asset(
|
|
file_path, project_id, folder_id, organization_id, metadata
|
|
)
|
|
|
|
if result["status"] == "success":
|
|
results["successful"] += 1
|
|
else:
|
|
results["failed"] += 1
|
|
results["errors"].append(
|
|
{"file": file_path, "error": result.get("message", "Unknown error")}
|
|
)
|
|
|
|
if results["failed"] > 0:
|
|
results["status"] = "partial"
|
|
|
|
results["message"] = (
|
|
f"Uploaded {results['successful']}/{results['total']} assets successfully"
|
|
)
|
|
|
|
return results
|
|
|
|
def test_connection(self) -> Dict:
|
|
"""
|
|
Test the API connection
|
|
|
|
Returns:
|
|
Dictionary containing connection test result
|
|
"""
|
|
try:
|
|
url = f"{self.base_url}/health"
|
|
|
|
response = requests.get(url, headers=self.headers, timeout=5)
|
|
|
|
if response.status_code == 200:
|
|
return {"status": "success", "message": "API connection successful"}
|
|
else:
|
|
return {
|
|
"status": "error",
|
|
"message": f"API returned status code {response.status_code}",
|
|
}
|
|
|
|
except Exception as e:
|
|
return {"status": "error", "message": f"Connection failed: {str(e)}"}
|