Underground-Digital/Workflow-Engine

1from typing import Optional2 3from pydantic import BaseModel, ConfigDict4 5from models.dataset import Document6from models.model import UploadFile7 8 9class NotionInfo(BaseModel):10    """11    Notion import info.12    """13 14    notion_workspace_id: str15    notion_obj_id: str16    notion_page_type: str17    document: Document = None18    tenant_id: str19    model_config = ConfigDict(arbitrary_types_allowed=True)20 21    def __init__(self, **data) -> None:22        super().__init__(**data)23 24 25class WebsiteInfo(BaseModel):26    """27    website import info.28    """29 30    provider: str31    job_id: str32    url: str33    mode: str34    tenant_id: str35    only_main_content: bool = False36 37    class Config:38        arbitrary_types_allowed = True39 40    def __init__(self, **data) -> None:41        super().__init__(**data)42 43 44class ExtractSetting(BaseModel):45    """46    Model class for provider response.47    """48 49    datasource_type: str50    upload_file: Optional[UploadFile] = None51    notion_info: Optional[NotionInfo] = None52    website_info: Optional[WebsiteInfo] = None53    document_model: Optional[str] = None54    model_config = ConfigDict(arbitrary_types_allowed=True)55 56    def __init__(self, **data) -> None:57        super().__init__(**data)58