"""Asset data models.
Defines AssetRecord for tracking generated media from ComfyUI workflows.
Uses stable identity (filename, subfolder, folder_type) instead of URLs.
"""
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
from urllib.parse import quote
@dataclass
class AssetRecord:
"""Record of a generated asset for tracking and viewing.
Uses (filename, subfolder, type) as stable identity instead of URL,
making the system robust to URL changes (e.g., different hostnames).
Attributes:
asset_id: UUID for tracking (ephemeral)
filename: Stable identity component (e.g., "ComfyUI_00123.png")
subfolder: Stable identity component (e.g., "" or "my_folder")
folder_type: Stable identity component (e.g., "output" or "temp")
prompt_id: Links to ComfyUI history
workflow_id: Workflow that generated this asset
created_at: Timestamp of creation
expires_at: TTL-based expiration timestamp
mime_type: Content type (e.g., "image/png")
width: Image width in pixels (if applicable)
height: Image height in pixels (if applicable)
bytes_size: File size in bytes
sha256: Content hash for deduplication (optional)
comfy_history: Full /history/{prompt_id} response for provenance
submitted_workflow: Original workflow submitted (for regeneration)
metadata: Extensible metadata dictionary
session_id: Conversation isolation identifier
"""
asset_id: str
filename: str
subfolder: str
folder_type: str
prompt_id: str
workflow_id: str
created_at: datetime
expires_at: datetime | None
# Presentation/display fields
mime_type: str
width: int | None
height: int | None
bytes_size: int
sha256: str | None # Content hash for deduplication
# Provenance (for regeneration and audit trail)
comfy_history: dict[str, Any] | None = field(default=None)
submitted_workflow: dict[str, Any] | None = field(default=None)
# Additional metadata
metadata: dict[str, Any] = field(default_factory=dict)
# Session tracking for conversation isolation
session_id: str | None = None
def get_asset_url(self, base_url: str) -> str:
"""Get asset URL for a given ComfyUI base URL.
Handles URL encoding for special characters in filenames and subfolders.
Also normalizes base_url (removes trailing slashes).
Args:
base_url: ComfyUI base URL (e.g., "http://localhost:8188")
Returns:
Full asset URL with proper encoding
"""
# Normalize base_url (remove trailing slash)
base_url = base_url.rstrip("/")
# URL encode filename and subfolder to handle special characters
encoded_filename = quote(self.filename, safe="")
encoded_subfolder = quote(self.subfolder, safe="") if self.subfolder else ""
# Build URL with proper encoding
if encoded_subfolder:
return f"{base_url}/view?filename={encoded_filename}&subfolder={encoded_subfolder}&type={self.folder_type}"
else:
return f"{base_url}/view?filename={encoded_filename}&type={self.folder_type}"
@property
def asset_url(self) -> str:
"""Compute asset URL on-the-fly from stable identity.
Note: This requires the base URL to be stored. The registry
sets this via _base_url attribute. Falls back to empty string
if base URL not available.
Returns:
Asset URL or empty string if base URL not set
"""
base_url = getattr(self, "_base_url", None)
if base_url:
return self.get_asset_url(base_url)
return ""
def set_base_url(self, base_url: str):
"""Set the ComfyUI base URL for computing asset URLs.
Args:
base_url: ComfyUI base URL
"""
self._base_url = base_url
