griptape-nodes 0.53.0__py3-none-any.whl → 0.54.1__py3-none-any.whl

griptape_nodes/cli/commands/models.py

@@ -0,0 +1,504 @@
+"""Models command for managing AI models."""
+
+import asyncio
+from typing import TYPE_CHECKING
+
+import typer
+from rich.table import Table
+
+from griptape_nodes.cli.shared import console
+from griptape_nodes.retained_mode.events.model_events import (
+    DeleteModelDownloadRequest,
+    DeleteModelDownloadResultFailure,
+    DeleteModelDownloadResultSuccess,
+    DeleteModelRequest,
+    DeleteModelResultFailure,
+    DeleteModelResultSuccess,
+    ListModelDownloadsRequest,
+    ListModelDownloadsResultFailure,
+    ListModelDownloadsResultSuccess,
+    ListModelsRequest,
+    ListModelsResultFailure,
+    ListModelsResultSuccess,
+    ModelInfo,
+    QueryInfo,
+    SearchModelsRequest,
+    SearchModelsResultFailure,
+    SearchModelsResultSuccess,
+)
+from griptape_nodes.retained_mode.retained_mode import GriptapeNodes
+
+if TYPE_CHECKING:
+    from griptape_nodes.retained_mode.events.model_events import ModelDownloadStatus
+
+app = typer.Typer(help="Manage AI models.")
+downloads_app = typer.Typer(help="Manage model download tracking records.")
+
+# Add downloads subcommand
+app.add_typer(downloads_app, name="downloads")
+
+
+@app.command("download")
+def download_command(
+    model_id: str = typer.Argument(..., help="Model ID or URL (e.g., 'microsoft/DialoGPT-medium')"),
+    local_dir: str | None = typer.Option(None, "--local-dir", help="Local directory to download the model to"),
+    revision: str = typer.Option("main", "--revision", help="Git revision to download"),
+) -> None:
+    """Download a model from Hugging Face Hub."""
+    asyncio.run(_download_model(model_id, local_dir, revision))
+
+
+@app.command("list")
+def list_command() -> None:
+    """List all downloaded model files in local cache."""
+    asyncio.run(_list_models())
+
+
+@app.command("delete")
+def delete_command(
+    model_id: str = typer.Argument(..., help="Model ID to delete (e.g., 'microsoft/DialoGPT-medium')"),
+) -> None:
+    """Delete model files from local cache."""
+    asyncio.run(_delete_model(model_id))
+
+
+@downloads_app.command("status")
+def downloads_status_command(
+    model_id: str = typer.Argument(None, help="Optional model ID to check download status for"),
+) -> None:
+    """Show download status for a specific model or all models."""
+    asyncio.run(_get_model_status(model_id))
+
+
+@downloads_app.command("list")
+def downloads_list_command() -> None:
+    """List all model download status records."""
+    asyncio.run(_get_model_status(None))
+
+
+@downloads_app.command("delete")
+def downloads_delete_command(
+    model_id: str = typer.Argument(
+        ..., help="Model ID to delete download status for (e.g., 'microsoft/DialoGPT-medium')"
+    ),
+) -> None:
+    """Delete download status tracking records for a model."""
+    asyncio.run(_delete_model_status(model_id))
+
+
+@app.command("search")
+def search_command(
+    query: str | None = typer.Argument(None, help="Search query to match against model names"),
+    task: str | None = typer.Option(None, "--task", help="Filter by task type"),
+    limit: int = typer.Option(20, "--limit", help="Maximum number of results (max: 100)"),
+    sort: str = typer.Option("downloads", "--sort", help="Sort results by"),
+    direction: str = typer.Option("desc", "--direction", help="Sort direction"),
+) -> None:
+    """Search for models on Hugging Face Hub."""
+    asyncio.run(_search_models(query, task, limit, sort, direction))
+
+
+async def _download_model(
+    model_id: str,
+    local_dir: str | None,
+    revision: str,
+) -> None:
+    """Download a model from Hugging Face Hub.
+
+    Args:
+        model_id: Model ID or URL to download
+        local_dir: Local directory to download the model to
+        revision: Git revision to download
+    """
+    console.print(f"[bold green]Downloading model: {model_id}[/bold green]")
+
+    try:
+        # ModelManager DownloadModelRequest will use this command so it's important that we don't use the request ourselves
+        model_manager = GriptapeNodes.ModelManager()
+        local_path = model_manager.download_model(
+            model_id=model_id,
+            local_dir=local_dir,
+            revision=revision,
+            allow_patterns=None,
+            ignore_patterns=None,
+        )
+
+        # Success case
+        console.print("[bold green]Model downloaded successfully![/bold green]")
+        console.print(f"[green]Downloaded to: {local_path}[/green]")
+
+    except Exception as e:
+        console.print("[bold red]Model download failed:[/bold red]")
+        console.print(f"[red]{e}[/red]")
+
+
+async def _list_models() -> None:
+    """List all downloaded models in the local cache."""
+    console.print("[bold green]Listing cached models...[/bold green]")
+
+    # Create the list request
+    request = ListModelsRequest()
+
+    try:
+        # Use the ModelManager to handle the listing
+        result = await GriptapeNodes.ahandle_request(request)
+        if isinstance(result, ListModelsResultSuccess):
+            # Success case
+            models = result.models
+            if models:
+                console.print(f"[bold green]Found {len(models)} cached models:[/bold green]")
+
+                table = Table()
+                table.add_column("Model ID", style="green")
+                table.add_column("Size (GB)", style="cyan", justify="right")
+
+                for model in models:
+                    size_gb = round((model.size_bytes or 0) / (1024**3), 2) if model.size_bytes else 0.0
+                    table.add_row(model.model_id, str(size_gb))
+                console.print(table)
+            else:
+                console.print("[yellow]No models found in local cache[/yellow]")
+
+        # Failure case
+
+        elif isinstance(result, ListModelsResultFailure):
+            console.print("[bold red]Model listing failed:[/bold red]")
+            if result.result_details:
+                console.print(f"[red]{result.result_details}[/red]")
+            if result.exception:
+                console.print(f"[dim]Error: {result.exception}[/dim]")
+        else:
+            console.print("[bold red]Model listing failed: Unknown error occurred[/bold red]")
+
+    except Exception as e:
+        console.print("[bold red]Unexpected error during model listing:[/bold red]")
+        console.print(f"[red]{e}[/red]")
+
+
+async def _delete_model(model_id: str) -> None:
+    """Delete a model from the local cache.
+
+    Args:
+        model_id: Model ID to delete
+    """
+    console.print(f"[bold yellow]Deleting model: {model_id}[/bold yellow]")
+
+    # Create the delete request
+    request = DeleteModelRequest(model_id=model_id)
+
+    try:
+        # Use the ModelManager to handle the deletion
+        result = await GriptapeNodes.ahandle_request(request)
+
+        if isinstance(result, DeleteModelResultSuccess):
+            # Success case
+            console.print("[bold green]Model deleted successfully![/bold green]")
+            console.print(f"[green]Deleted: {result.deleted_path}[/green]")
+        # Failure case
+
+        elif isinstance(result, DeleteModelResultFailure):
+            console.print("[bold red]Model deletion failed:[/bold red]")
+            if result.result_details:
+                console.print(f"[red]{result.result_details}[/red]")
+            if result.exception:
+                console.print(f"[dim]Error: {result.exception}[/dim]")
+        else:
+            console.print("[bold red]Model deletion failed: Unknown error occurred[/bold red]")
+
+    except Exception as e:
+        console.print("[bold red]Unexpected error during model deletion:[/bold red]")
+        console.print(f"[red]{e}[/red]")
+
+
+def _format_download_row(download: "ModelDownloadStatus") -> tuple[str, str, str, str, str, str]:
+    """Format a download status object into table row data.
+
+    Args:
+        download: ModelDownloadStatus object
+
+    Returns:
+        tuple: (model_id, status_colored, progress_str, size_str, eta_str, started_str)
+    """
+    progress_str = _format_progress(download)
+    size_str = _format_size(download)
+    eta_str = _format_eta(download)
+    started_str = _format_timestamp(download)
+    status_colored = _format_status(download)
+
+    return (
+        download.model_id,
+        status_colored,
+        progress_str,
+        size_str,
+        eta_str,
+        started_str,
+    )
+
+
+def _format_progress(download: "ModelDownloadStatus") -> str:
+    """Format download progress information."""
+    if download.total_files is not None and download.completed_files is not None and download.total_files > 0:
+        progress_percent = (download.completed_files / download.total_files) * 100
+        return f"{progress_percent:.1f}%"
+    return "Unknown"
+
+
+def _format_size(download: "ModelDownloadStatus") -> str:
+    """Format download size information."""
+    if download.total_files is not None and download.completed_files is not None:
+        return f"{download.completed_files}/{download.total_files} files"
+    return "Unknown"
+
+
+def _format_eta(download: "ModelDownloadStatus") -> str:
+    """Format estimated time of arrival."""
+    # ETA is not available in the current ModelDownloadStatus structure
+    # For active downloads, we could potentially calculate based on progress
+    # but without timing data, we return a status-appropriate message
+    if download.status == "downloading":
+        return "In progress"
+    if download.status == "completed":
+        return "Completed"
+    if download.status == "failed":
+        return "Failed"
+    return "Unknown"
+
+
+def _format_timestamp(download: "ModelDownloadStatus") -> str:
+    """Format started timestamp."""
+    started_at = download.started_at
+    if not started_at:
+        return "Unknown"
+
+    try:
+        from datetime import datetime
+
+        dt = datetime.fromisoformat(started_at)
+        return dt.strftime("%H:%M:%S")
+    except Exception:
+        return started_at[:10]  # Fallback
+
+
+def _format_status(download: "ModelDownloadStatus") -> str:
+    """Format status with color coding."""
+    status = download.status
+    status_colors = {
+        "completed": "green",
+        "failed": "red",
+        "downloading": "yellow",
+    }
+
+    if status in status_colors:
+        return f"[{status_colors[status]}]{status}[/{status_colors[status]}]"
+    return status
+
+
+def _display_downloads_table(downloads: list["ModelDownloadStatus"]) -> None:
+    """Display downloads in a formatted table.
+
+    Args:
+        downloads: List of ModelDownloadStatus objects
+    """
+    console.print(f"[bold green]Found {len(downloads)} download(s):[/bold green]")
+
+    table = Table()
+    table.add_column("Model ID", style="green")
+    table.add_column("Status", style="cyan")
+    table.add_column("Progress", style="yellow", justify="right")
+    table.add_column("Size", style="blue", justify="right")
+    table.add_column("ETA", style="magenta", justify="right")
+    table.add_column("Started", style="dim")
+
+    for download in downloads:
+        row_data = _format_download_row(download)
+        table.add_row(*row_data)
+
+    console.print(table)
+
+
+async def _get_model_status(model_id: str | None) -> None:
+    """Get download status for models.
+
+    Args:
+        model_id: Optional model ID to get status for
+    """
+    if model_id:
+        console.print(f"[bold green]Getting download status for: {model_id}[/bold green]")
+    else:
+        console.print("[bold green]Getting download status for all models...[/bold green]")
+
+    # Create the status request
+    request = ListModelDownloadsRequest(model_id=model_id)
+
+    try:
+        # Use the ModelManager to handle the status query
+        result = await GriptapeNodes.ahandle_request(request)
+
+        if isinstance(result, ListModelDownloadsResultSuccess):
+            # Success case
+            downloads = result.downloads
+            if downloads:
+                _display_downloads_table(downloads)
+            elif model_id:
+                console.print(f"[yellow]No download found for model: {model_id}[/yellow]")
+            else:
+                console.print("[yellow]No downloads found[/yellow]")
+
+        elif isinstance(result, ListModelDownloadsResultFailure):
+            console.print("[bold red]Failed to get download status:[/bold red]")
+            if result.result_details:
+                console.print(f"[red]{result.result_details}[/red]")
+            if result.exception:
+                console.print(f"[dim]Error: {result.exception}[/dim]")
+        else:
+            console.print("[bold red]Failed to get download status: Unknown error occurred[/bold red]")
+
+    except Exception as e:
+        console.print("[bold red]Unexpected error getting download status:[/bold red]")
+        console.print(f"[red]{e}[/red]")
+
+
+async def _search_models(
+    query: str | None,
+    task: str | None,
+    limit: int,
+    sort: str,
+    direction: str,
+) -> None:
+    """Search for models on Hugging Face Hub.
+
+    Args:
+        query: Search query to match against model names
+        task: Filter by task type
+        limit: Maximum number of results
+        sort: Sort results by
+        direction: Sort direction
+    """
+    if query:
+        console.print(f"[bold green]Searching for models: {query}[/bold green]")
+    else:
+        console.print("[bold green]Searching for models...[/bold green]")
+
+    # Create the search request
+    request = SearchModelsRequest(
+        query=query,
+        task=task,
+        library=None,
+        author=None,
+        tags=None,
+        limit=limit,
+        sort=sort,
+        direction=direction,
+    )
+
+    try:
+        # Use the ModelManager to handle the search
+        result = await GriptapeNodes.ahandle_request(request)
+
+        if isinstance(result, SearchModelsResultSuccess):
+            # Success case
+            models = result.models
+            if models:
+                _display_search_results(models, result.query_info)
+            else:
+                console.print("[yellow]No models found matching the search criteria[/yellow]")
+
+        elif isinstance(result, SearchModelsResultFailure):
+            console.print("[bold red]Model search failed:[/bold red]")
+            if result.result_details:
+                console.print(f"[red]{result.result_details}[/red]")
+            if result.exception:
+                console.print(f"[dim]Error: {result.exception}[/dim]")
+        else:
+            console.print("[bold red]Model search failed: Unknown error occurred[/bold red]")
+
+    except Exception as e:
+        console.print("[bold red]Unexpected error during model search:[/bold red]")
+        console.print(f"[red]{e}[/red]")
+
+
+def _display_search_results(models: list[ModelInfo], query_info: QueryInfo) -> None:
+    """Display model search results in a formatted table.
+
+    Args:
+        models: List of model information
+        query_info: Information about the search query
+    """
+    console.print(f"[bold green]Found {len(models)} models[/bold green]")
+
+    # Show search parameters if any were used
+    params = []
+    if query_info.query:
+        params.append(f"query: {query_info.query}")
+    if query_info.task:
+        params.append(f"task: {query_info.task}")
+    if query_info.library:
+        params.append(f"library: {query_info.library}")
+    if query_info.author:
+        params.append(f"author: {query_info.author}")
+    if query_info.tags:
+        params.append(f"tags: {', '.join(query_info.tags)}")
+
+    if params:
+        console.print(f"[dim]Search parameters: {', '.join(params)}[/dim]")
+
+    table = Table()
+    table.add_column("Model ID", style="green")
+    table.add_column("Author", style="blue")
+    table.add_column("Downloads", style="cyan", justify="right")
+    table.add_column("Likes", style="yellow", justify="right")
+    table.add_column("Task", style="magenta")
+    table.add_column("Library", style="dim")
+
+    for model in models:
+        downloads_str = f"{model.downloads:,}" if model.downloads else "0"
+        likes_str = str(model.likes or 0)
+        task_str = model.task or ""
+        library_str = model.library or ""
+        author_str = model.author or ""
+
+        table.add_row(
+            model.model_id,
+            author_str,
+            downloads_str,
+            likes_str,
+            task_str,
+            library_str,
+        )
+
+    console.print(table)
+
+
+async def _delete_model_status(model_id: str) -> None:
+    """Delete download status records for a model.
+
+    Args:
+        model_id: Model ID to delete download status for
+    """
+    console.print(f"[bold yellow]Deleting download status for: {model_id}[/bold yellow]")
+
+    # Create the delete request
+    request = DeleteModelDownloadRequest(model_id=model_id)
+
+    try:
+        # Use the ModelManager to handle the deletion
+        result = await GriptapeNodes.ahandle_request(request)
+
+        if isinstance(result, DeleteModelDownloadResultSuccess):
+            # Success case
+            console.print("[bold green]Download status deleted successfully![/bold green]")
+            console.print(f"[green]Deleted status file: {result.deleted_path}[/green]")
+        # Failure case
+
+        elif isinstance(result, DeleteModelDownloadResultFailure):
+            console.print("[bold red]Download status deletion failed:[/bold red]")
+            if result.result_details:
+                console.print(f"[red]{result.result_details}[/red]")
+            if result.exception:
+                console.print(f"[dim]Error: {result.exception}[/dim]")
+        else:
+            console.print("[bold red]Download status deletion failed: Unknown error occurred[/bold red]")
+
+    except Exception as e:
+        console.print("[bold red]Unexpected error during download status deletion:[/bold red]")
+        console.print(f"[red]{e}[/red]")

griptape_nodes/cli/commands/self.py

@@ -11,10 +11,9 @@ from griptape_nodes.cli.shared import (
     GITHUB_UPDATE_URL,
     LATEST_TAG,
     PYPI_UPDATE_URL,
-    config_manager,
     console,
-    os_manager,
 )
+from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
 from griptape_nodes.utils.uv_utils import find_uv_bin
 from griptape_nodes.utils.version_utils import (
     get_complete_version_string,
@@ -23,6 +22,10 @@ from griptape_nodes.utils.version_utils import (
     get_latest_version_pypi,
 )
 
+config_manager = GriptapeNodes.ConfigManager()
+secrets_manager = GriptapeNodes.SecretsManager()
+os_manager = GriptapeNodes.OSManager()
+
 app = typer.Typer(help="Manage this CLI installation.")
 
 

griptape_nodes/cli/main.py

@@ -9,7 +9,8 @@ from rich.console import Console
 # Add current directory to path for imports to work
 sys.path.append(str(Path.cwd()))
 
-from griptape_nodes.cli.commands import config, engine, init, libraries, self
+from griptape_nodes.cli.commands import config, engine, init, libraries, models, self
+from griptape_nodes.utils.version_utils import get_complete_version_string
 
 console = Console()
 
@@ -26,17 +27,26 @@ app.command("init", help="Initialize engine configuration.")(init.init_command)
 app.add_typer(config.app, name="config")
 app.add_typer(self.app, name="self")
 app.add_typer(libraries.app, name="libraries")
+app.add_typer(models.app, name="models")
 app.command("engine")(engine.engine_command)
 
 
 @app.callback()
 def main(
     ctx: typer.Context,
+    version: bool = typer.Option(  # noqa: FBT001
+        False, "--version", help="Show version and exit."
+    ),
     no_update: bool = typer.Option(  # noqa: FBT001
         False, "--no-update", help="Skip the auto-update check."
     ),
 ) -> None:
     """Griptape Nodes Engine CLI."""
+    if version:
+        version_string = get_complete_version_string()
+        console.print(f"[bold green]{version_string}[/bold green]")
+        raise typer.Exit
+
     if ctx.invoked_subcommand is None:
         # Default to engine command when no subcommand is specified
         engine.engine_command(no_update=no_update)

griptape_nodes/cli/shared.py

@@ -8,10 +8,6 @@ from typing import Any
 from rich.console import Console
 from xdg_base_dirs import xdg_config_home, xdg_data_home
 
-from griptape_nodes.retained_mode.managers.config_manager import ConfigManager
-from griptape_nodes.retained_mode.managers.os_manager import OSManager
-from griptape_nodes.retained_mode.managers.secrets_manager import SecretsManager
-
 
 @dataclass
 class InitConfig:
@@ -61,11 +57,6 @@ ENV_LIBRARIES_SYNC = (
 ENV_GTN_BUCKET_NAME = os.getenv("GTN_BUCKET_NAME")
 ENV_LIBRARIES_BASE_DIR = os.getenv("GTN_LIBRARIES_BASE_DIR", str(DATA_DIR / "libraries"))
 
-# Initialize managers
-config_manager = ConfigManager()
-secrets_manager = SecretsManager(config_manager)
-os_manager = OSManager()
-
 
 def init_system_config() -> None:
     """Initializes the system config directory if it doesn't exist."""

griptape_nodes/common/directed_graph.py

@@ -12,6 +12,10 @@ class DirectedGraph:
         self._nodes: set[str] = set()
         self._predecessors: dict[str, set[str]] = {}
 
+    def __len__(self) -> int:
+        """Return the number of nodes in the graph."""
+        return len(self._nodes)
+
     def add_node(self, node_for_adding: str) -> None:
         """Add a node to the graph."""
         self._nodes.add(node_for_adding)
@@ -31,9 +35,21 @@ class DirectedGraph:
     def in_degree(self, node: str) -> int:
         """Return the in-degree of a node (number of incoming edges)."""
         if node not in self._nodes:
-            return 0
+            msg = f"Node {node} not found in graph"
+            raise KeyError(msg)
         return len(self._predecessors.get(node, set()))
 
+    def out_degree(self, node: str) -> int:
+        """Return the out-degree of a node (number of outgoing edges)."""
+        if node not in self._nodes:
+            msg = f"Node {node} not found in graph"
+            raise KeyError(msg)
+        count = 0
+        for predecessors in self._predecessors.values():
+            if node in predecessors:
+                count += 1
+        return count
+
     def remove_node(self, node: str) -> None:
         """Remove a node and all its edges from the graph."""
         if node not in self._nodes: