Source code for bigdata_client.api.uploads

import datetime
from typing import Iterable, Literal, Optional, Protocol

from pydantic import BaseModel, Field, model_validator

from bigdata_client.daterange import AbsoluteDateRange
from bigdata_client.enum_utils import StrEnum
from bigdata_client.models.sharing import SharePermission




[docs]
class FileStatus(StrEnum):
    PENDING = "PENDING"
    PROCESSING = "PROCESSING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    DELETED = "DELETED"






[docs]
class FileResponse(BaseModel):
    """Model to represent a single upload in the ListFilesResponse"""

    id: str = Field(alias="file_id")
    name: str = Field(alias="file_name")
    status: FileStatus
    uploaded_at: datetime.datetime = Field(alias="upload_ts")
    raw_size: int
    folder_id: Optional[str] = None
    tags: Optional[list[str]] = None
    company_shared_permission: Optional[SharePermission] = None



[docs]
    @model_validator(mode="before")
    @classmethod
    def set_company_shared_permission(cls, values: dict):
        if values.get("shared_with"):
            values["company_shared_permission"] = (
                SharePermission.READ
            )  # We only share with own organization
        return values





[docs]
    @model_validator(mode="after")
    def set_tz(self):
        self.uploaded_at = self.uploaded_at.replace(tzinfo=datetime.timezone.utc)

        return self








[docs]
class ListFilesResponse(BaseModel):
    """Model to represent the response of the list uploads endpoint"""

    results: list[FileResponse]






[docs]
class ExtractorTypes(StrEnum):
    PDF_EXTRACTOR_1_0 = "PDF_EXTRACTOR_1_0"






[docs]
class PostFileRequest(BaseModel):
    """Model to represent the request of the post upload endpoint"""

    filename: str
    folder_id: Optional[str] = None
    source_url: Optional[str] = None
    upload_mode: Optional[str] = None
    properties: Optional[dict] = None
    connect_to: list[Literal["smart_topics"]] = ["smart_topics"]
    extractor: Optional[ExtractorTypes] = None






[docs]
class PostFileResponse(BaseModel):
    """Model to represent the response of the post upload endpoint"""

    file_id: str
    location: str = Field(alias="Location")






[docs]
class GetFileStatusResponse(BaseModel):
    status: Optional[FileStatus] = None
    error: Optional[str] = None






[docs]
class DeleteFileResponse(BaseModel):
    message: str






[docs]
class GetDownloadPresignedUrlResponse(BaseModel):
    url: str






[docs]
class UploadsConnectionProtocol(Protocol):


[docs]
    def list_files(


        self,
        date_range: Optional[AbsoluteDateRange],
        tags: Optional[list[str]],
        status: Optional[FileStatus],
        file_name: Optional[str],
        folder_id: Optional[str],
        page_size: int,
        shared: Optional[bool],
        offset: int,
    ) -> ListFilesResponse: ...



[docs]
    def get_file(self, id: str) -> FileResponse: ...





[docs]
    def post_file(self, request: PostFileRequest) -> PostFileResponse: ...





[docs]
    def get_file_status(self, file_id: str) -> GetFileStatusResponse: ...





[docs]
    def delete_file(self, id: str) -> DeleteFileResponse: ...





[docs]
    def get_download_presigned_url(


        self, id: str
    ) -> GetDownloadPresignedUrlResponse: ...



[docs]
    def download_analytics(self, id: str) -> Iterable[bytes]: ...





[docs]
    def download_annotated(self, id: str) -> Iterable[bytes]: ...





[docs]
    def download_text(self, id: str) -> Iterable[bytes]: ...





[docs]
    def share_file_with_company(self, file_id: str) -> FileResponse: ...





[docs]
    def unshare_file_with_company(self, file_id: str) -> FileResponse: ...





[docs]
    def update_file_tags(self, file_id: str, tags: list[str]) -> FileResponse: ...








[docs]
class FileShareResponse(BaseModel):
    status: Literal["OK"]
    message: str