sql_qa_pair/app/endpoints/schema.py

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class CardFilter(BaseModel):
    """Metadata describing a filter bound to a BI card."""

    card_id: str
    filter_id: str
    type: str
    where_clause: str
    default_value: Optional[str] = None
    options: List[str] = Field(default_factory=list)


class CardInfo(BaseModel):
    card_id: str
    card_name: str
    card_desc: Optional[str] = None
    dataset_id: str
    sql_select: str
    sql_where: Optional[Dict[str, Any]] = None
    sql_groupby: Optional[str] = None
    filters: List[CardFilter] = Field(default_factory=list)


class DashboardInfo(BaseModel):
    dashboard_id: str
    dashboard_name: str
    dashboard_desc: Optional[str] = None
    folder_path: Optional[str] = None
    cards: List[CardInfo] = Field(default_factory=list)


class QARequest(BaseModel):
    request_id: str
    dashboard_id: str
    card_ids: List[str]
    bbk_id: str
    user_token: Optional[str] = None
    user_request: Optional[str] = None


class QAPair(BaseModel):
    dashboard_id: str
    dashboard_name: str
    dashboard_desc: Optional[str]
    card_id: str
    card_name: str
    card_desc: Optional[str]
    dataset_id: str
    question: str
    answer: str = ""
    question_with_slot: str = ""
    filter_ids: List[str] = Field(default_factory=list)
    filters: List[str] = Field(default_factory=list)
    defs: List[str] = Field(default_factory=list)


class QAResponseBody(BaseModel):
    dashboard_id: str
    bbk_id: str
    card_ids: List[str]
    qa_pairs: List[QAPair]


class QAResponse(BaseModel):
    returnCode: str = "SUCCESS"
    body: QAResponseBody