- 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
27 lines
644 B
Python
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]
|