Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 41 additions & 15 deletions src/madengine/cli/commands/discover.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,14 @@ def discover(
List[str],
typer.Option("--tags", "-t", help="Model tags to discover (can specify multiple)"),
] = [],
full: Annotated[
bool,
typer.Option("--full", "-f", help="Output full JSON with all discovered models and their tags"),
] = False,
json: Annotated[
bool,
typer.Option("--json", "-j", help="Output plain JSON only (no formatting, no status messages)"),
] = False,
verbose: Annotated[
bool, typer.Option("--verbose", "-v", help="Enable verbose logging")
] = False,
Expand All @@ -40,35 +48,53 @@ def discover(
limits selection to models under ``scripts/<scope>/`` (e.g.
``MAD-private/inference`` → models named ``MAD-private/...`` with tag
``inference``). Use ``scope/all`` for every model in that scope.

**Full JSON output** (``--full``): outputs complete model cards with all tags
and metadata in JSON format, similar to ``--tags`` output but for all models.

**Plain JSON output** (``--json``): outputs only pure JSON without any formatting
or status messages. Useful for piping to other tools or CI/CD pipelines.
"""
setup_logging(verbose)
# Skip logging setup if json mode is enabled
if not json:
setup_logging(verbose)

# Process tags to handle comma-separated values
processed_tags = split_comma_separated_tags(tags)

console.print(
Panel(
f"🔍 [bold cyan]Discovering Models[/bold cyan]\n"
f"Tags: [yellow]{processed_tags if processed_tags else 'All models'}[/yellow]",
title="Model Discovery",
border_style="blue",
# Skip console output if json mode is enabled
if not json:
console.print(
Panel(
f"🔍 [bold cyan]Discovering Models[/bold cyan]\n"
f"Tags: [yellow]{processed_tags if processed_tags else 'All models'}[/yellow]\n"
f"Full JSON: [yellow]{full}[/yellow]",
title="Model Discovery",
border_style="blue",
)
)
)

try:
# Create args namespace similar to mad.py
args = create_args_namespace(tags=processed_tags)
args = create_args_namespace(tags=processed_tags, full=full, json=json)

# Use DiscoverModels class
# Note: DiscoverModels prints output directly and returns None
discover_models_instance = DiscoverModels(args=args)
result = discover_models_instance.run()

console.print("✅ [bold green]Model discovery completed successfully[/bold green]")

# Skip success message if json mode is enabled
if not json:
console.print("✅ [bold green]Model discovery completed successfully[/bold green]")

except Exception as e:
console.print(f"💥 [bold red]Model discovery failed: {e}[/bold red]")
if verbose:
console.print_exception()
# In json mode, write errors to stderr to keep stdout clean
if json:
import sys
print(f"Error: {e}", file=sys.stderr)
else:
console.print(f"💥 [bold red]Model discovery failed: {e}[/bold red]")
if verbose:
console.print_exception()
raise typer.Exit(ExitCode.FAILURE)

38 changes: 34 additions & 4 deletions src/madengine/utils/discover_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -448,15 +448,45 @@ def select_models(self) -> None:
self.selected_models.extend(tag_models)

def print_models(self) -> None:
# Check if --full flag is set to output full JSON with all models
output_full_json = getattr(self.args, 'full', False)
# Check if --json flag is set for plain JSON output without formatting
json_mode = getattr(self.args, 'json', False)

if self.selected_models:
# print selected models using parsed tags and adding backslash-separated extra args
self.rich_console.print(f"[bold green]📋 Selected Models ({len(self.selected_models)} models):[/bold green]")
if not json_mode:
self.rich_console.print(f"[bold green]📋 Selected Models ({len(self.selected_models)} models):[/bold green]")
print(json.dumps(self.selected_models, indent=4))
elif output_full_json:
# Output full JSON with all discovered models and their complete model cards
# Include both regular models and expanded custom models
all_models = self.models.copy()

# Expand and add custom models
for custom_model in self.custom_models:
custom_model.update_model()
dirname = custom_model.name.split("/")[0]
custom_model.dockerfile = os.path.normpath(
os.path.join("scripts", dirname, custom_model.dockerfile)
)
custom_model.scripts = os.path.normpath(
os.path.join("scripts", dirname, custom_model.scripts)
)
all_models.append(custom_model.to_dict())

if not json_mode:
self.rich_console.print(f"[bold cyan]📊 All Models with Full Details ({len(all_models)} models):[/bold cyan]")
print(json.dumps(all_models, indent=4))
else:
# print list of all model names
self.rich_console.print(f"[bold cyan]📊 Available Models ({len(self.model_list)} total):[/bold cyan]")
for model_name in self.model_list:
print(f" {model_name}")
if not json_mode:
self.rich_console.print(f"[bold cyan]📊 Available Models ({len(self.model_list)} total):[/bold cyan]")
for model_name in self.model_list:
print(f" {model_name}")
else:
# In json mode without --full or --tags, output model names as JSON array
print(json.dumps(self.model_list, indent=4))

def run(self, live_output: bool = True):

Expand Down