comfygit-core 0.2.0__py3-none-any.whl

comfygit_core/services/model_downloader.py

@@ -0,0 +1,422 @@
+"""Model download service for fetching models from URLs."""
+
+from __future__ import annotations
+
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+
+import requests
+from blake3 import blake3
+
+from ..configs.model_config import ModelConfig
+from ..logging.logging_config import get_logger
+from ..models.exceptions import DownloadErrorContext
+from ..models.shared import ModelWithLocation
+from ..utils.model_categories import get_model_category
+
+if TYPE_CHECKING:
+    from ..repositories.model_repository import ModelRepository
+    from ..repositories.workspace_config_repository import WorkspaceConfigRepository
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class DownloadRequest:
+    """Request to download a model."""
+    url: str
+    target_path: Path  # Full path in global models directory
+    workflow_name: str | None = None
+
+
+@dataclass
+class DownloadResult:
+    """Result of a download operation."""
+    success: bool
+    model: ModelWithLocation | None = None
+    error: str | None = None
+    error_context: "DownloadErrorContext | None" = None  # Structured error info
+
+
+class ModelDownloader:
+    """Handles model downloads with hashing and indexing.
+
+    Responsibilities:
+    - Download files from URLs with progress tracking
+    - Compute hashes (short + full blake3)
+    - Register in ModelRepository
+    - Detect URL type (civitai/HF/direct)
+    """
+
+    def __init__(
+        self,
+        model_repository: ModelRepository,
+        workspace_config: WorkspaceConfigRepository,
+        models_dir: Path | None = None
+    ):
+        """Initialize ModelDownloader.
+
+        Args:
+            model_repository: Repository for indexing models
+            workspace_config: Workspace config for API credentials and models directory
+            models_dir: Optional override for models directory (defaults to workspace config)
+        """
+        self.repository = model_repository
+        self.workspace_config = workspace_config
+
+        # Use provided models_dir or get from workspace config
+        self.models_dir = models_dir if models_dir is not None else workspace_config.get_models_directory()
+
+        # Since workspace always has models_dir configured, this should never be None
+        # Raise clear error if it somehow is
+        if self.models_dir is None:
+            raise ValueError(
+                "No models directory available. Either provide models_dir parameter "
+                "or ensure workspace config has a models directory configured."
+            )
+
+        self.model_config = ModelConfig.load()
+
+    def detect_url_type(self, url: str) -> str:
+        """Detect source type from URL.
+
+        Args:
+            url: URL to analyze
+
+        Returns:
+            'civitai', 'huggingface', or 'custom'
+        """
+        url_lower = url.lower()
+
+        if "civitai.com" in url_lower:
+            return "civitai"
+        elif "huggingface.co" in url_lower or "hf.co" in url_lower:
+            return "huggingface"
+        else:
+            return "custom"
+
+    def suggest_path(
+        self,
+        url: str,
+        node_type: str | None = None,
+        filename_hint: str | None = None
+    ) -> Path:
+        """Suggest download path based on context.
+
+        For known nodes: checkpoints/model.safetensors
+        For unknown: Uses filename hint or extracts from URL
+
+        Args:
+            url: Download URL
+            node_type: Optional node type for category mapping
+            filename_hint: Optional filename hint from workflow
+
+        Returns:
+            Suggested relative path (including base directory)
+        """
+        # Extract filename from URL or use hint
+        filename = self._extract_filename(url, filename_hint)
+
+        # If node type is known, map to directory
+        if node_type and self.model_config.is_model_loader_node(node_type):
+            directories = self.model_config.get_directories_for_node(node_type)
+            base_dir = directories[0]  # e.g., "checkpoints"
+            return Path(base_dir) / filename
+
+        # Fallback: try to extract category from filename hint
+        if filename_hint:
+            category = get_model_category(filename_hint)
+            return Path(category) / filename
+
+        # Last resort: use generic models directory
+        return Path("models") / filename
+
+    def _extract_filename(self, url: str, filename_hint: str | None = None) -> str:
+        """Extract filename from URL or use hint.
+
+        Args:
+            url: Download URL
+            filename_hint: Optional filename from workflow
+
+        Returns:
+            Filename to use
+        """
+        # Try to extract from URL path
+        parsed = urlparse(url)
+        url_filename = Path(parsed.path).name
+
+        # Use URL filename if it looks valid (has extension)
+        if url_filename and '.' in url_filename:
+            return url_filename
+
+        # Fall back to hint
+        if filename_hint:
+            # Extract just the filename from hint path
+            return Path(filename_hint).name
+
+        # Last resort: generate generic name
+        return "downloaded_model.safetensors"
+
+    def _check_provider_auth(self, provider: str) -> bool:
+        """Check if authentication is configured for a provider.
+
+        Args:
+            provider: Provider type ('civitai', 'huggingface', 'custom')
+
+        Returns:
+            True if auth credentials are configured
+        """
+        if provider == "civitai":
+            if not self.workspace_config:
+                return False
+            api_key = self.workspace_config.get_civitai_token()
+            return api_key is not None and api_key.strip() != ""
+        elif provider == "huggingface":
+            # Check HF_TOKEN environment variable
+            import os
+            token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+            return token is not None and token.strip() != ""
+        else:
+            return False
+
+    def _classify_download_error(
+        self,
+        error: Exception,
+        url: str,
+        provider: str,
+        has_auth: bool
+    ) -> DownloadErrorContext:
+        """Classify download error and create structured context.
+
+        Args:
+            error: The exception that occurred
+            url: Download URL
+            provider: Provider type
+            has_auth: Whether auth was configured
+
+        Returns:
+            DownloadErrorContext with classification
+        """
+        from urllib.error import URLError
+        from socket import timeout as SocketTimeout
+
+        http_status = None
+        error_category = "unknown"
+        raw_error = str(error)
+
+        # Classify based on exception type
+        if isinstance(error, requests.HTTPError):
+            http_status = error.response.status_code
+
+            if http_status == 401:
+                # Unauthorized - check if we have auth
+                if not has_auth:
+                    error_category = "auth_missing"
+                else:
+                    error_category = "auth_invalid"
+            elif http_status == 403:
+                # Forbidden - could be rate limit, permissions, or invalid token
+                if not has_auth and provider in ("civitai", "huggingface"):
+                    error_category = "auth_missing"
+                else:
+                    error_category = "forbidden"
+            elif http_status == 404:
+                error_category = "not_found"
+            elif http_status >= 500:
+                error_category = "server"
+            else:
+                error_category = "unknown"
+
+        elif isinstance(error, (URLError, SocketTimeout, requests.Timeout, requests.ConnectionError)):
+            error_category = "network"
+
+        return DownloadErrorContext(
+            provider=provider,
+            error_category=error_category,
+            http_status=http_status,
+            url=url,
+            has_configured_auth=has_auth,
+            raw_error=raw_error
+        )
+
+    def download(
+        self,
+        request: DownloadRequest,
+        progress_callback=None
+    ) -> DownloadResult:
+        """Download and index a model.
+
+        Flow:
+        1. Check if URL already downloaded
+        2. Validate URL and target path
+        3. Download to temp file with progress
+        4. Hash during download (streaming)
+        5. Move to target location
+        6. Register in repository
+        7. Add source URL
+
+        Args:
+            request: Download request with URL and target path
+            progress_callback: Optional callback(bytes_downloaded, total_bytes) for progress updates.
+                             total_bytes may be None if server doesn't provide Content-Length.
+
+        Returns:
+            DownloadResult with model or error
+        """
+        temp_path: Path | None = None
+        try:
+            # Step 1: Check if already downloaded
+            existing = self.repository.find_by_source_url(request.url)
+            if existing:
+                logger.info(f"Model already downloaded from URL: {existing.relative_path}")
+                return DownloadResult(success=True, model=existing)
+
+            # Step 2: Validate target path
+            request.target_path.parent.mkdir(parents=True, exist_ok=True)
+
+            # Step 3-4: Download with streaming hash calculation
+            logger.info(f"Downloading from {request.url}")
+
+            # Add Civitai auth header if URL is from Civitai and we have an API key
+            headers = {}
+            if "civitai.com" in request.url.lower() and self.workspace_config:
+                api_key = self.workspace_config.get_civitai_token()
+                if api_key:
+                    headers['Authorization'] = f'Bearer {api_key}'
+                    logger.debug("Using Civitai API key for authentication")
+
+            # Timeout: (connect_timeout, read_timeout)
+            # 30s to establish connection, None for read (allow slow downloads)
+            response = requests.get(request.url, stream=True, timeout=(30, None), headers=headers)
+            response.raise_for_status()
+
+            # Extract total size from headers (may be None)
+            total_size = None
+            if 'content-length' in response.headers:
+                try:
+                    total_size = int(response.headers['content-length'])
+                except (ValueError, TypeError):
+                    pass
+
+            # Use temp file for atomic move
+            with tempfile.NamedTemporaryFile(delete=False, dir=request.target_path.parent) as temp_file:
+                temp_path = Path(temp_file.name)
+
+                # Stream download with hash calculation
+                hasher = blake3()
+                file_size = 0
+
+                chunk_size = 8192
+                for chunk in response.iter_content(chunk_size=chunk_size):
+                    if chunk:
+                        temp_file.write(chunk)
+                        hasher.update(chunk)
+                        file_size += len(chunk)
+
+                        if progress_callback:
+                            progress_callback(file_size, total_size)
+
+            # Step 5: Calculate short hash for indexing
+            short_hash = self.repository.calculate_short_hash(temp_path)
+            blake3_hash = hasher.hexdigest()
+
+            # Step 6: Atomic move to final location (replace handles existing files)
+            temp_path.replace(request.target_path)
+            temp_path = None  # Clear temp_path since file has been moved
+
+            # Step 7: Register in repository
+            relative_path = request.target_path.relative_to(self.models_dir)
+            mtime = request.target_path.stat().st_mtime
+
+            self.repository.ensure_model(
+                hash=short_hash,
+                file_size=file_size,
+                blake3_hash=blake3_hash
+            )
+
+            self.repository.add_location(
+                model_hash=short_hash,
+                base_directory=self.models_dir,
+                relative_path=str(relative_path),
+                filename=request.target_path.name,
+                mtime=mtime
+            )
+
+            # Step 8: Add source URL
+            source_type = self.detect_url_type(request.url)
+            self.repository.add_source(
+                model_hash=short_hash,
+                source_type=source_type,
+                source_url=request.url
+            )
+
+            # Step 9: Create result model
+            model = ModelWithLocation(
+                hash=short_hash,
+                file_size=file_size,
+                blake3_hash=blake3_hash,
+                sha256_hash=None,
+                relative_path=str(relative_path),
+                filename=request.target_path.name,
+                mtime=mtime,
+                last_seen=int(mtime),
+                metadata={}
+            )
+
+            logger.info(f"Successfully downloaded and indexed: {relative_path}")
+            return DownloadResult(success=True, model=model)
+
+        except requests.HTTPError as e:
+            # HTTP errors with status codes - classify them
+            provider = self.detect_url_type(request.url)
+            has_auth = self._check_provider_auth(provider)
+            error_context = self._classify_download_error(e, request.url, provider, has_auth)
+
+            # Generate user-friendly message
+            user_message = error_context.get_user_message()
+            logger.error(f"Download failed: {user_message}")
+
+            return DownloadResult(
+                success=False,
+                error=user_message,
+                error_context=error_context
+            )
+
+        except (requests.Timeout, requests.ConnectionError) as e:
+            # Network errors
+            provider = self.detect_url_type(request.url)
+            error_context = self._classify_download_error(e, request.url, provider, False)
+            user_message = error_context.get_user_message()
+            logger.error(f"Download failed: {user_message}")
+
+            return DownloadResult(
+                success=False,
+                error=user_message,
+                error_context=error_context
+            )
+
+        except Exception as e:
+            # Unexpected errors - still provide some context
+            provider = self.detect_url_type(request.url)
+            has_auth = self._check_provider_auth(provider)
+            error_context = self._classify_download_error(e, request.url, provider, has_auth)
+            user_message = error_context.get_user_message()
+            logger.error(f"Unexpected download error: {user_message}")
+
+            return DownloadResult(
+                success=False,
+                error=user_message,
+                error_context=error_context
+            )
+
+        finally:
+            # Always clean up temp file if it still exists (download failed or was interrupted)
+            if temp_path is not None and temp_path.exists():
+                try:
+                    temp_path.unlink()
+                    logger.debug(f"Cleaned up temporary file: {temp_path}")
+                except Exception as cleanup_error:
+                    logger.warning(f"Failed to clean up temp file {temp_path}: {cleanup_error}")

comfygit_core/services/node_lookup_service.py

@@ -0,0 +1,251 @@
+"""NodeLookupService - Pure stateless service for finding nodes and analyzing requirements."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from comfygit_core.models.exceptions import CDNodeNotFoundError, CDRegistryError
+from comfygit_core.models.shared import NodeInfo
+
+from ..analyzers.custom_node_scanner import CustomNodeScanner
+from ..caching import APICacheManager, CustomNodeCacheManager
+from ..clients import ComfyRegistryClient, GitHubClient
+from ..logging.logging_config import get_logger
+from ..utils.git import is_git_url
+
+if TYPE_CHECKING:
+    from comfygit_core.repositories.node_mappings_repository import NodeMappingsRepository
+    from comfygit_core.repositories.workspace_config_repository import WorkspaceConfigRepository
+
+logger = get_logger(__name__)
+
+
+class NodeLookupService:
+    """Pure stateless service for finding nodes and analyzing their requirements.
+
+    Responsibilities:
+    - Registry lookup (cache-first, then API)
+    - GitHub API calls (validating repos, getting commit info)
+    - Requirement scanning (analyzing node directories)
+    - Cache management (API responses, downloaded node archives)
+    """
+
+    def __init__(
+        self,
+        cache_path: Path,
+        node_mappings_repository: NodeMappingsRepository | None = None,
+        workspace_config_repository: WorkspaceConfigRepository | None = None,
+    ):
+        """Initialize the node lookup service.
+
+        Args:
+            cache_path: Required path to workspace cache directory
+            node_mappings_repository: Repository for cached node mappings
+            workspace_config_repository: Repository for workspace config (cache preference)
+        """
+        self.scanner = CustomNodeScanner()
+        self.api_cache = APICacheManager(cache_base_path=cache_path)
+        self.custom_node_cache = CustomNodeCacheManager(cache_base_path=cache_path)
+        self.registry_client = ComfyRegistryClient(cache_manager=self.api_cache)
+        self.github_client = GitHubClient(cache_manager=self.api_cache)
+        self.node_mappings_repository = node_mappings_repository
+        self.workspace_config_repository = workspace_config_repository
+
+    def find_node(self, identifier: str) -> NodeInfo | None:
+        """Find node info from cache, registry, or git URL.
+
+        Cache-first strategy:
+        1. If prefer_registry_cache=True and package in local cache → return from cache
+        2. Otherwise → query live API
+
+        Args:
+            identifier: Registry ID, git URL, or name. Supports @version/@ref syntax:
+                       - registry-id@1.0.0 (registry version)
+                       - https://github.com/user/repo@v1.2.3 (git tag)
+                       - https://github.com/user/repo@main (git branch)
+                       - https://github.com/user/repo@abc123 (git commit)
+
+        Returns:
+            NodeInfo with metadata, or None if not found
+        """
+        # Parse version/ref from identifier (e.g., "package-id@1.2.3" or "https://...@branch")
+        requested_version = None
+        base_identifier = identifier
+
+        if '@' in identifier:
+            parts = identifier.rsplit('@', 1)  # rsplit to handle URLs with @
+            base_identifier = parts[0]
+            requested_version = parts[1]
+
+        # Check if it's a git URL - these bypass cache
+        if is_git_url(base_identifier):
+            try:
+                if repo_info := self.github_client.get_repository_info(base_identifier, ref=requested_version):
+                    return NodeInfo(
+                        name=repo_info.name,
+                        repository=repo_info.clone_url,
+                        source="git",
+                        version=repo_info.latest_commit  # This will be the requested ref's commit
+                    )
+            except Exception as e:
+                logger.warning(f"Invalid git URL: {e}")
+                return None
+
+        # Check if we should prefer cached mappings
+        prefer_cache = True
+        if self.workspace_config_repository:
+            prefer_cache = self.workspace_config_repository.get_prefer_registry_cache()
+
+        # Strategy: Cache first, then API
+        if prefer_cache and self.node_mappings_repository:
+            package = self.node_mappings_repository.get_package(base_identifier)
+            if package:
+                logger.debug(f"Found '{base_identifier}' in local cache")
+                return NodeInfo.from_global_package(package, version=requested_version)
+            else:
+                logger.debug(f"'{base_identifier}' not in local cache, trying API...")
+
+        # Fallback to registry API
+        try:
+            registry_node = self.registry_client.get_node(base_identifier)
+            if registry_node:
+                if requested_version:
+                    version = requested_version
+                    logger.debug(f"Using requested version: {version}")
+                else:
+                    version = registry_node.latest_version.version if registry_node.latest_version else None
+                node_version = self.registry_client.install_node(registry_node.id, version)
+                if node_version:
+                    registry_node.latest_version = node_version
+                return NodeInfo.from_registry_node(registry_node)
+        except CDRegistryError as e:
+            logger.warning(f"Cannot reach registry: {e}")
+
+        logger.debug(f"Node '{base_identifier}' not found")
+        return None
+
+    def get_node(self, identifier: str) -> NodeInfo:
+        """Get a node - raises if not found.
+
+        Args:
+            identifier: Registry ID, node name, or git URL
+
+        Returns:
+            NodeInfo with metadata
+
+        Raises:
+            CDNodeNotFoundError: If node not found in any source
+        """
+        node = self.find_node(identifier)
+        if not node:
+            # Build context-aware error based on what was tried
+            if is_git_url(identifier):
+                msg = f"Node '{identifier}' not found. GitHub repository is invalid or inaccessible."
+            else:
+                # Registry lookup was attempted
+                prefer_cache = True
+                if self.workspace_config_repository:
+                    prefer_cache = self.workspace_config_repository.get_prefer_registry_cache()
+
+                sources_tried = ["local cache", "registry API"] if prefer_cache else ["registry API"]
+                msg = f"Node '{identifier}' not found in {' or '.join(sources_tried)}"
+
+            raise CDNodeNotFoundError(msg)
+        return node
+
+    def scan_requirements(self, node_path: Path) -> list[str]:
+        """Scan a node directory for Python requirements.
+
+        Args:
+            node_path: Path to node directory
+
+        Returns:
+            List of requirement strings (empty if none found)
+        """
+        deps = self.scanner.scan_node(node_path)
+        if deps and deps.requirements:
+            logger.info(f"Found {len(deps.requirements)} requirements in {node_path.name}")
+            return deps.requirements
+        logger.info(f"No requirements found in {node_path.name}")
+        return []
+
+    def download_to_cache(self, node_info: NodeInfo) -> Path | None:
+        """Download a node to cache and return the cached path.
+
+        Args:
+            node_info: Node information
+
+        Returns:
+            Path to cached node directory, or None if download failed
+        """
+        import tempfile
+
+        from ..utils.download import download_and_extract_archive
+        from ..utils.git import git_clone
+
+        # Check if already cached
+        if cache_path := self.custom_node_cache.get_cached_path(node_info):
+            logger.debug(f"Node '{node_info.name}' already in cache")
+            return cache_path
+
+        # Download to temp location
+        with tempfile.TemporaryDirectory() as tmpdir:
+            temp_path = Path(tmpdir) / node_info.name
+
+            try:
+                if node_info.source == "registry":
+                    if not node_info.download_url:
+                        # Fallback: Clone from repository if download URL missing
+                        if node_info.repository:
+                            logger.info(
+                                f"No CDN package for '{node_info.name}', "
+                                f"falling back to git clone from {node_info.repository}"
+                            )
+                            # Update source to git for this installation
+                            node_info.source = "git"
+                            ref = node_info.version if node_info.version else None
+                            git_clone(node_info.repository, temp_path, depth=1, ref=ref, timeout=30)
+                        else:
+                            logger.error(
+                                f"Cannot download '{node_info.name}': "
+                                f"no CDN package and no repository URL"
+                            )
+                            return None
+                    else:
+                        download_and_extract_archive(node_info.download_url, temp_path)
+                elif node_info.source == "git":
+                    if not node_info.repository:
+                        logger.error(f"No repository URL for git node '{node_info.name}'")
+                        return None
+                    ref = node_info.version if node_info.version else None
+                    git_clone(node_info.repository, temp_path, depth=1, ref=ref, timeout=30)
+                else:
+                    logger.error(f"Unsupported source: '{node_info.source}'")
+                    return None
+
+                # Cache it
+                logger.info(f"Caching node '{node_info.name}'")
+                return self.custom_node_cache.cache_node(node_info, temp_path)
+
+            except Exception as e:
+                logger.error(f"Failed to download node '{node_info.name}': {e}")
+                return None
+
+    def search_nodes(self, query: str, limit: int = 10) -> list[NodeInfo] | None:
+        """Search for nodes in the registry.
+
+        Args:
+            query: Search term
+            limit: Maximum results
+
+        Returns:
+            List of matching NodeInfo objects or None
+        """
+        try:
+            nodes = self.registry_client.search_nodes(query)
+            if nodes:
+                return [NodeInfo.from_registry_node(node) for node in nodes[:limit]]
+        except CDRegistryError as e:
+            logger.warning(f"Failed to search registry: {e}")
+        return None