label_ai_service/app/models/qa_models.py

from pydantic import BaseModel


class TextQAItem(BaseModel):
    subject: str
    predicate: str
    object: str
    source_snippet: str


class GenTextQARequest(BaseModel):
    items: list[TextQAItem]
    model: str | None = None
    prompt_template: str | None = None


class QAPair(BaseModel):
    question: str
    answer: str


class ImageQAItem(BaseModel):
    subject: str
    predicate: str
    object: str
    qualifier: str | None = None
    cropped_image_path: str


class GenImageQARequest(BaseModel):
    items: list[ImageQAItem]
    model: str | None = None
    prompt_template: str | None = None


class ImageQAPair(BaseModel):
    question: str
    answer: str
    image_path: str


class TextQAResponse(BaseModel):
    pairs: list[QAPair]


class ImageQAResponse(BaseModel):
    pairs: list[ImageQAPair]