velai/services/text/TextGenerator.py

from abc import ABC
from dataclasses import dataclass, field
from typing import List

from PIL.Image import Image

from velai.services.generator_service import BaseGeneratorService, GenerationInput, GenerationResult


@dataclass
class TextGenerationInput(GenerationInput):
    prompt: str = ""
    images: List[Image] = field(default_factory=list)
    reasoning: bool | None = None
    temperature: float | None = None
    max_tokens: int | None = None


@dataclass
class TextGenerationResult(GenerationResult):
    """Result of a text generation call."""

    text: str = ""


class TextGenerator(BaseGeneratorService[TextGenerationInput, TextGenerationResult], ABC):
    pass