hokusai/bulkgen/providers/models.py
Konstantin Fickel d0dac5b1bf
Some checks failed
Continuous Integration / Build Package (push) Successful in 30s
Continuous Integration / Lint, Check & Test (push) Failing after 38s
refactor: move model definitions into providers and extract resolve module
- Rename ImageProvider to BlackForestProvider, TextProvider to MistralProvider
- Add get_provided_models() abstract method to Provider base class
- Move model lists from models.py into each provider's get_provided_models()
- Add providers/registry.py to aggregate models from all providers
- Extract infer_required_capabilities and resolve_model from config.py to resolve.py
- Update tests to use new names and import paths
2026-02-15 11:03:57 +01:00

27 lines
644 B
Python

"""Model types and capability definitions for AI providers."""
from __future__ import annotations
from dataclasses import dataclass
from enum import StrEnum
from typing import Literal
class Capability(StrEnum):
"""Capabilities a model may support."""
TEXT_GENERATION = "text generation"
VISION = "vision"
TEXT_TO_IMAGE = "text-to-image"
REFERENCE_IMAGES = "reference images"
CONTROL_IMAGES = "control images"
@dataclass(frozen=True)
class ModelInfo:
"""Describes a supported model and its capabilities."""
name: str
provider: str
type: Literal["text", "image"]
capabilities: list[Capability]