Sylvester Meighan
e14ffc2284
- Add BrainApiClient for interacting with Brain API storage endpoints - Update MediaStreamInput to use storage_id + filename instead of presigned URLs - Update MediaStreamOutput to use Brain API client for uploads and remove output_object_keys dependency - Add environment configuration for Brain API connection - Add test file for Brain API client functionality This enables ComfyUI nodes to work with the new storage architecture where Brain API acts as a proxy for MinIO operations.
287 lines
10 KiB
Python
287 lines
10 KiB
Python
"""
|
|
Brain API Client for ComfyUI Nodes
|
|
|
|
This client provides methods to interact with the Brain API storage endpoints,
|
|
replacing the need for pre-signed URLs in the ComfyUI workflow.
|
|
"""
|
|
|
|
import requests
|
|
import os
|
|
import logging
|
|
from typing import Optional, Dict, Any
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
dotenv_path = os.path.join(current_dir, ".env")
|
|
load_dotenv(dotenv_path=dotenv_path)
|
|
|
|
# Setup logging
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
|
|
|
|
class BrainApiClient:
|
|
"""
|
|
Client for interacting with Brain API storage endpoints.
|
|
|
|
This client handles authentication and provides methods for uploading,
|
|
downloading, and deleting files through the Brain API storage endpoints.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the Brain API client with configuration from environment variables."""
|
|
self.base_url = os.getenv("BRANDO_BRAIN_API_BASE_URL", "http://localhost:2024/api")
|
|
self.api_key = os.getenv("BRANDO_API_KEY")
|
|
|
|
if not self.api_key:
|
|
raise ValueError(
|
|
"BRANDO_API_KEY environment variable is required for Brain API authentication"
|
|
)
|
|
|
|
self.headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"User-Agent": "ComfyUI-NilorNodes/1.0"
|
|
}
|
|
|
|
logging.info(f"Brain API Client initialized with base URL: {self.base_url}")
|
|
|
|
def upload_file_to_storage(self, file_path: str, filename: str) -> Dict[str, Any]:
|
|
"""
|
|
Upload a file to Brain API storage and return storage metadata.
|
|
|
|
Args:
|
|
file_path: Local path to the file to upload
|
|
filename: Name to use for the uploaded file
|
|
|
|
Returns:
|
|
Dict containing storage_id and filename
|
|
|
|
Raises:
|
|
requests.RequestException: If upload fails
|
|
FileNotFoundError: If file_path doesn't exist
|
|
"""
|
|
if not os.path.exists(file_path):
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
|
|
url = f"{self.base_url}/storage/upload"
|
|
|
|
try:
|
|
with open(file_path, 'rb') as file:
|
|
files = {'file': (filename, file, 'application/octet-stream')}
|
|
|
|
logging.info(f"Uploading file '{filename}' to Brain API storage...")
|
|
response = requests.post(
|
|
url,
|
|
files=files,
|
|
headers=self.headers,
|
|
timeout=300
|
|
)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
logging.info(f"Upload successful. Storage ID: {result.get('storage_id')}")
|
|
return result
|
|
|
|
except requests.RequestException as e:
|
|
logging.error(f"Failed to upload file '{filename}': {e}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error uploading file '{filename}': {e}")
|
|
raise
|
|
|
|
def upload_fileobj_to_storage(self, file_obj, filename: str, content_type: str = 'application/octet-stream') -> Dict[str, Any]:
|
|
"""
|
|
Upload a file-like object to Brain API storage and return storage metadata.
|
|
|
|
Args:
|
|
file_obj: File-like object to upload
|
|
filename: Name to use for the uploaded file
|
|
content_type: MIME type of the file
|
|
|
|
Returns:
|
|
Dict containing storage_id and filename
|
|
|
|
Raises:
|
|
requests.RequestException: If upload fails
|
|
"""
|
|
url = f"{self.base_url}/storage/upload"
|
|
|
|
try:
|
|
files = {'file': (filename, file_obj, content_type)}
|
|
|
|
logging.info(f"Uploading file object '{filename}' to Brain API storage...")
|
|
response = requests.post(
|
|
url,
|
|
files=files,
|
|
headers=self.headers,
|
|
timeout=300
|
|
)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
logging.info(f"Upload successful. Storage ID: {result.get('storage_id')}")
|
|
return result
|
|
|
|
except requests.RequestException as e:
|
|
logging.error(f"Failed to upload file object '{filename}': {e}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error uploading file object '{filename}': {e}")
|
|
raise
|
|
|
|
def download_file_from_storage(self, storage_id: str, filename: str, dest_path: str) -> str:
|
|
"""
|
|
Download a file from Brain API storage to a local path.
|
|
|
|
Args:
|
|
storage_id: Storage ID of the file to download
|
|
filename: Name of the file to download
|
|
dest_path: Local path where the file should be saved
|
|
|
|
Returns:
|
|
Path to the downloaded file
|
|
|
|
Raises:
|
|
requests.RequestException: If download fails
|
|
"""
|
|
url = f"{self.base_url}/storage/{storage_id}"
|
|
params = {'filename': filename}
|
|
|
|
try:
|
|
logging.info(f"Downloading file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
|
response = requests.get(
|
|
url,
|
|
params=params,
|
|
headers=self.headers,
|
|
timeout=300,
|
|
stream=True
|
|
)
|
|
response.raise_for_status()
|
|
|
|
# Ensure destination directory exists
|
|
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
|
|
|
with open(dest_path, 'wb') as f:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
|
|
logging.info(f"Download successful. File saved to: {dest_path}")
|
|
return dest_path
|
|
|
|
except requests.RequestException as e:
|
|
logging.error(f"Failed to download file '{filename}' (storage_id: {storage_id}): {e}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error downloading file '{filename}': {e}")
|
|
raise
|
|
|
|
def get_file_from_storage(self, storage_id: str, filename: str) -> bytes:
|
|
"""
|
|
Get file content from Brain API storage as bytes.
|
|
|
|
Args:
|
|
storage_id: Storage ID of the file to download
|
|
filename: Name of the file to download
|
|
|
|
Returns:
|
|
File content as bytes
|
|
|
|
Raises:
|
|
requests.RequestException: If download fails
|
|
"""
|
|
url = f"{self.base_url}/storage/{storage_id}"
|
|
params = {'filename': filename}
|
|
|
|
try:
|
|
logging.info(f"Getting file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
|
response = requests.get(
|
|
url,
|
|
params=params,
|
|
headers=self.headers,
|
|
timeout=300
|
|
)
|
|
response.raise_for_status()
|
|
|
|
logging.info(f"File retrieval successful. Size: {len(response.content)} bytes")
|
|
return response.content
|
|
|
|
except requests.RequestException as e:
|
|
logging.error(f"Failed to get file '{filename}' (storage_id: {storage_id}): {e}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error getting file '{filename}': {e}")
|
|
raise
|
|
|
|
def delete_file_from_storage(self, storage_id: str, filename: str) -> None:
|
|
"""
|
|
Delete a file from Brain API storage.
|
|
|
|
Args:
|
|
storage_id: Storage ID of the file to delete
|
|
filename: Name of the file to delete
|
|
|
|
Raises:
|
|
requests.RequestException: If deletion fails
|
|
"""
|
|
url = f"{self.base_url}/storage/{storage_id}"
|
|
params = {'filename': filename}
|
|
|
|
try:
|
|
logging.info(f"Deleting file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
|
response = requests.delete(
|
|
url,
|
|
params=params,
|
|
headers=self.headers,
|
|
timeout=60
|
|
)
|
|
response.raise_for_status()
|
|
|
|
logging.info(f"File deletion successful")
|
|
|
|
except requests.RequestException as e:
|
|
logging.error(f"Failed to delete file '{filename}' (storage_id: {storage_id}): {e}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error deleting file '{filename}': {e}")
|
|
raise
|
|
|
|
def health_check(self) -> bool:
|
|
"""
|
|
Check if the Brain API is accessible and authentication is working.
|
|
|
|
Returns:
|
|
True if API is accessible, False otherwise
|
|
"""
|
|
try:
|
|
# Try to access a simple endpoint to verify connectivity
|
|
url = f"{self.base_url}/health" # Assuming there's a health endpoint
|
|
response = requests.get(url, headers=self.headers, timeout=10)
|
|
return response.status_code == 200
|
|
except:
|
|
# If health endpoint doesn't exist, try the storage upload endpoint
|
|
# with a HEAD request to check authentication
|
|
try:
|
|
url = f"{self.base_url}/storage/upload"
|
|
response = requests.head(url, headers=self.headers, timeout=10)
|
|
return response.status_code in [200, 405] # 405 Method Not Allowed is OK for HEAD
|
|
except:
|
|
return False
|
|
|
|
|
|
# Global client instance
|
|
_brain_api_client = None
|
|
|
|
def get_brain_api_client() -> BrainApiClient:
|
|
"""
|
|
Get or create the global Brain API client instance.
|
|
|
|
Returns:
|
|
BrainApiClient instance
|
|
"""
|
|
global _brain_api_client
|
|
if _brain_api_client is None:
|
|
_brain_api_client = BrainApiClient()
|
|
return _brain_api_client
|