feat: add bulkgen models command listing available models with capabilities
Some checks failed
Continuous Integration / Build Package (push) Successful in 37s
Continuous Integration / Lint, Check & Test (push) Failing after 47s

This commit is contained in:
Konstantin Fickel 2026-02-14 23:06:17 +01:00
parent 47b53db760
commit b536ff9d79
Signed by: kfickel
GPG key ID: A793722F9933C1A5
2 changed files with 160 additions and 0 deletions

View file

@ -13,6 +13,7 @@ import typer
from bulkgen.builder import BuildEvent, BuildResult, run_build from bulkgen.builder import BuildEvent, BuildResult, run_build
from bulkgen.config import ProjectConfig, load_config from bulkgen.config import ProjectConfig, load_config
from bulkgen.graph import build_graph, get_build_order from bulkgen.graph import build_graph, get_build_order
from bulkgen.models import ALL_MODELS
app = typer.Typer(name="bulkgen", help="AI artifact build tool.") app = typer.Typer(name="bulkgen", help="AI artifact build tool.")
@ -176,3 +177,53 @@ def graph() -> None:
if preds: if preds:
arrow = click.style(" <- ", dim=True) arrow = click.style(" <- ", dim=True)
click.echo(f" {node}{arrow}{', '.join(preds)}") click.echo(f" {node}{arrow}{', '.join(preds)}")
@app.command()
def models() -> None:
"""List available models and their capabilities."""
name_width = max(len(m.name) for m in ALL_MODELS)
provider_width = max(len(m.provider) for m in ALL_MODELS)
type_width = max(len(m.type) for m in ALL_MODELS)
header_name = "Model".ljust(name_width)
header_provider = "Provider".ljust(provider_width)
header_type = "Type".ljust(type_width)
header_caps = "Capabilities"
click.echo(
click.style(header_name, bold=True)
+ " "
+ click.style(header_provider, bold=True)
+ " "
+ click.style(header_type, bold=True)
+ " "
+ click.style(header_caps, bold=True)
)
click.echo(
"" * name_width
+ " "
+ "" * provider_width
+ " "
+ "" * type_width
+ " "
+ "" * len(header_caps)
)
for model in ALL_MODELS:
name_col = model.name.ljust(name_width)
provider_col = model.provider.ljust(provider_width)
type_col = model.type.ljust(type_width)
caps_col = ", ".join(model.capabilities)
type_color = "green" if model.type == "image" else "cyan"
click.echo(
click.style(name_col, fg="white", bold=True)
+ " "
+ provider_col
+ " "
+ click.style(type_col, fg=type_color)
+ " "
+ click.style(caps_col, dim=True)
)

109
bulkgen/models.py Normal file
View file

@ -0,0 +1,109 @@
"""Registry of supported models and their capabilities."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal
@dataclass(frozen=True)
class ModelInfo:
"""Describes a supported model and its capabilities."""
name: str
provider: str
type: Literal["text", "image"]
capabilities: list[str]
TEXT_MODELS: list[ModelInfo] = [
ModelInfo(
name="mistral-large-latest",
provider="Mistral",
type="text",
capabilities=["text generation"],
),
ModelInfo(
name="mistral-small-latest",
provider="Mistral",
type="text",
capabilities=["text generation"],
),
ModelInfo(
name="pixtral-large-latest",
provider="Mistral",
type="text",
capabilities=["text generation", "vision"],
),
ModelInfo(
name="pixtral-12b-latest",
provider="Mistral",
type="text",
capabilities=["text generation", "vision"],
),
]
IMAGE_MODELS: list[ModelInfo] = [
ModelInfo(
name="flux-dev",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image"],
),
ModelInfo(
name="flux-pro",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image"],
),
ModelInfo(
name="flux-pro-1.1",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image"],
),
ModelInfo(
name="flux-pro-1.1-ultra",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "high resolution"],
),
ModelInfo(
name="flux-2-pro",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "reference images"],
),
ModelInfo(
name="flux-kontext-pro",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "reference images"],
),
ModelInfo(
name="flux-pro-1.0-canny",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "control images", "edge detection"],
),
ModelInfo(
name="flux-pro-1.0-depth",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "control images", "depth map"],
),
ModelInfo(
name="flux-pro-1.0-fill",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "inpainting"],
),
ModelInfo(
name="flux-pro-1.0-expand",
provider="BlackForestLabs",
type="image",
capabilities=["text-to-image", "outpainting"],
),
]
ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS