refactor: move model definitions into providers and extract resolve module
Some checks failed
Continuous Integration / Build Package (push) Successful in 30s
Continuous Integration / Lint, Check & Test (push) Failing after 38s

- 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
This commit is contained in:
Konstantin Fickel 2026-02-15 11:03:57 +01:00
parent dc6a75f5c4
commit d0dac5b1bf
Signed by: kfickel
GPG key ID: A793722F9933C1A5
13 changed files with 432 additions and 390 deletions

View file

@ -1,4 +1,4 @@
"""Registry of supported models and their capabilities."""
"""Model types and capability definitions for AI providers."""
from __future__ import annotations
@ -25,106 +25,3 @@ class ModelInfo:
provider: str
type: Literal["text", "image"]
capabilities: list[Capability]
TEXT_MODELS: list[ModelInfo] = [
ModelInfo(
name="mistral-large-latest",
provider="Mistral",
type="text",
capabilities=[Capability.TEXT_GENERATION],
),
ModelInfo(
name="mistral-small-latest",
provider="Mistral",
type="text",
capabilities=[Capability.TEXT_GENERATION],
),
ModelInfo(
name="pixtral-large-latest",
provider="Mistral",
type="text",
capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
),
ModelInfo(
name="pixtral-12b-latest",
provider="Mistral",
type="text",
capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
),
]
IMAGE_MODELS: list[ModelInfo] = [
ModelInfo(
name="flux-dev",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE],
),
ModelInfo(
name="flux-pro",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE],
),
ModelInfo(
name="flux-pro-1.1",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE],
),
ModelInfo(
name="flux-pro-1.1-ultra",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE],
),
ModelInfo(
name="flux-2-pro",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
),
ModelInfo(
name="flux-kontext-pro",
provider="BlackForestLabs",
type="image",
capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
),
ModelInfo(
name="flux-pro-1.0-canny",
provider="BlackForestLabs",
type="image",
capabilities=[
Capability.TEXT_TO_IMAGE,
Capability.CONTROL_IMAGES,
],
),
ModelInfo(
name="flux-pro-1.0-depth",
provider="BlackForestLabs",
type="image",
capabilities=[
Capability.TEXT_TO_IMAGE,
Capability.CONTROL_IMAGES,
],
),
ModelInfo(
name="flux-pro-1.0-fill",
provider="BlackForestLabs",
type="image",
capabilities=[
Capability.TEXT_TO_IMAGE,
],
),
ModelInfo(
name="flux-pro-1.0-expand",
provider="BlackForestLabs",
type="image",
capabilities=[
Capability.TEXT_TO_IMAGE,
],
),
]
ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS