backend/src/api/services/order_export.py - master - FlysAt/DataFoundry

backend/src/api/services/order_export.py
473 строки19 KB
AlexanderShmygol
fixed
17 май 2026, 14:28
17 май 2026, 14:28446906b
import csv
import json
import math
from io import BytesIO
from io import StringIO
from pathlib import PurePosixPath
from urllib.parse import urlsplit, urlunsplit
import zipfile

from fastapi import HTTPException, Request, status
from sqlalchemy.ext.asyncio import AsyncSession

from src.api.repositories.label_studio_annotation_current import LabelStudioAnnotationCurrentRepository
from src.api.repositories.label_studio_project import LabelStudioProjectRepository
from src.api.repositories.label_studio_task import LabelStudioTaskRepository
from src.api.repositories.order_annotation_type import OrderAnnotationTypeRepository
from src.api.repositories.order_annotation_type_executor import OrderAnnotationTypeExecutorRepository
from src.api.repositories.order import OrderRepository
from src.api.repositories.order_dataset import OrderDatasetRepository
from src.api.storage.minio import MinioStorage
from src.core.config import minio_settings
from src.models.order_dataset import DatasetPart, OrderDataset
from src.models.order import Order, OrderStatus
from src.models.user import User

TEXT_DATASET_TAG = "text"


class OrderExportService:
    def __init__(self, db: AsyncSession) -> None:
        self.orders = OrderRepository(db)
        self.datasets = OrderDatasetRepository(db)
        self.annotation_types = OrderAnnotationTypeRepository(db)
        self.executor_annotation_types = OrderAnnotationTypeExecutorRepository(db)
        self.projects = LabelStudioProjectRepository(db)
        self.tasks = LabelStudioTaskRepository(db)
        self.annotations = LabelStudioAnnotationCurrentRepository(db)
        self.storage = MinioStorage()

    async def build_annotated_dataset_archive(self, order_id: int, current_user: User) -> tuple[bytes, str]:
        order, dataset = await self._get_order_with_dataset(order_id, current_user)
        if order.status != OrderStatus.COMPLETED:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Annotated dataset can be downloaded only after order is completed",
            )
        if self._is_text_order(order):
            return await self._build_annotated_text_dataset_json(order, dataset)

        current_parts = await self.datasets.list_current_parts(dataset.id)
        dataset_objects = self._validated_objects(current_parts)
        labeling_objects = await self._current_labeling_objects(dataset.id, dataset_objects)
        labeling_prefix = f"orders/{order.id}/datasets/{dataset.id}/labeling/"

        if not dataset_objects and not labeling_objects:
            raise HTTPException(status_code=404, detail="No exported dataset files found")

        archive_buffer = BytesIO()
        used_archive_paths: set[str] = set()
        used_archive_filenames: dict[str, set[str]] = {}

        with zipfile.ZipFile(archive_buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
            for item in dataset_objects:
                relative_path = PurePosixPath(item["relative_path"])
                archive_path = self._unique_archive_path(
                    PurePosixPath("dataset") / relative_path,
                    used_archive_paths,
                    used_archive_filenames,
                )
                archive.writestr(
                    archive_path,
                    self.storage.download_bytes(item["object_key"], bucket_name=item["bucket"]),
                )

            for object_key in labeling_objects:
                relative_path = PurePosixPath(object_key.removeprefix(labeling_prefix))
                flattened_name = relative_path.name
                if len(relative_path.parts) > 1:
                    flattened_name = "_".join(relative_path.parts[:-1] + (relative_path.name,))
                archive_path = self._unique_archive_path(
                    PurePosixPath("labeling") / flattened_name,
                    used_archive_paths,
                    used_archive_filenames,
                )
                archive.writestr(
                    archive_path,
                    self.storage.download_bytes(object_key),
                )

        archive_buffer.seek(0)
        filename = f"annotated_dataset_{order.id}.zip"
        return archive_buffer.getvalue(), filename

    async def _build_annotated_text_dataset_json(self, order: Order, dataset: OrderDataset) -> tuple[bytes, str]:
        rows = await self._annotated_text_rows(order, dataset)
        payload = json.dumps(rows, ensure_ascii=False, indent=2).encode("utf-8")
        return payload, f"annotated_dataset_{order.id}.json"

    async def build_annotated_dataset_preview(
        self,
        order_id: int,
        current_user: User,
        request: Request,
    ) -> list[dict] | tuple[bytes, str]:
        order, dataset = await self._get_order_with_dataset(order_id, current_user)
        if self._is_text_order(order):
            rows = await self._annotated_text_rows(order, dataset)
            preview_count = max(1, math.ceil(len(rows) * 0.2))
            payload = json.dumps(rows[:preview_count], ensure_ascii=False, indent=2).encode("utf-8")
            return payload, f"annotated_dataset_{order.id}_preview.json"

        annotation_types = await self._list_label_annotation_names(order.id)
        if not annotation_types:
            raise HTTPException(status_code=400, detail="Order annotation types are not configured")
        annotation_class_ids = {
            annotation: index
            for index, annotation in enumerate(annotation_types)
        }

        project = await self.projects.get_by_order_dataset(dataset.id)
        if not project:
            raise HTTPException(status_code=404, detail="Labeling project not found")

        current_objects = self._validated_objects(await self.datasets.list_current_parts(dataset.id))
        current_buckets_by_key = {
            item["object_key"]: item["bucket"]
            for item in current_objects
        }

        tasks = await self.tasks.list_by_project_id(project.id)
        tasks = [
            task
            for task in tasks
            if task.source_object_key in current_buckets_by_key
        ]
        if not tasks:
            raise HTTPException(status_code=404, detail="Annotated dataset preview is not available")

        annotations = await self.annotations.list_by_project_id(project.id)
        preview_annotations_by_task_id = {
            annotation.task_id: self._extract_preview_annotations(annotation.result_payload, annotation_class_ids)
            for annotation in annotations
        }
        annotated_tasks = [
            task
            for task in tasks
            if preview_annotations_by_task_id.get(task.id)
        ]

        if not annotated_tasks:
            raise HTTPException(status_code=404, detail="Annotated dataset preview is not available")

        preview_count = max(1, math.ceil(len(tasks) * 0.2))
        preview_tasks = annotated_tasks[:preview_count]
        if len(preview_tasks) < preview_count:
            preview_task_ids = {task.id for task in preview_tasks}
            preview_tasks.extend(
                task
                for task in tasks
                if task.id not in preview_task_ids
            )
            preview_tasks = preview_tasks[:preview_count]

        preview_items: list[dict] = []
        public_minio_endpoint = self._resolve_public_url(
            minio_settings.MINIO_PUBLIC_ENDPOINT,
            request,
        )

        for task in preview_tasks:
            preview_items.append(
                {
                    "id": task.id,
                    "image_url": self.storage.get_presigned_get_url(
                        task.source_object_key,
                        bucket_name=current_buckets_by_key[task.source_object_key],
                        public_endpoint=public_minio_endpoint,
                    ),
                    "annotation_source": "txt",
                    "annotations": preview_annotations_by_task_id.get(task.id, []),
                }
            )

        return preview_items

    async def _annotated_text_rows(self, order: Order, dataset: OrderDataset) -> list[dict]:
        project = await self.projects.get_by_order_dataset(dataset.id)
        if not project:
            raise HTTPException(status_code=404, detail="Labeling project not found")

        current_rows = self._current_text_rows(await self.datasets.list_current_parts(dataset.id))
        if not current_rows:
            raise HTTPException(status_code=404, detail="Validated text dataset rows not found")

        tasks = [
            task
            for task in await self.tasks.list_by_project_id(project.id)
            if task.source_object_key in current_rows
        ]
        annotations_by_task_id = {
            annotation.task_id: annotation
            for annotation in await self.annotations.list_by_project_id(project.id)
        }

        rows: list[dict] = []
        for task in tasks:
            annotation = annotations_by_task_id.get(task.id)
            if annotation is None:
                continue
            sentiment = self._extract_text_sentiment(annotation.result_payload)
            if sentiment is None:
                continue

            rows.append(
                {
                    "text": current_rows[task.source_object_key],
                    "id": task.external_task_id,
                    "sentiment": sentiment,
                    "annotator": self._extract_annotator(annotation.result_payload, annotation.completed_by),
                    "annotation_id": annotation.external_annotation_id,
                    "created_at": self._payload_timestamp(annotation.result_payload, "created_at", annotation.created_at),
                    "updated_at": self._payload_timestamp(annotation.result_payload, "updated_at", annotation.updated_at),
                    "lead_time": annotation.result_payload.get("lead_time"),
                }
            )

        if not rows:
            raise HTTPException(status_code=404, detail="Annotated text dataset is not available")
        return rows

    async def _get_order_with_dataset(self, order_id: int, current_user: User) -> tuple[Order, OrderDataset]:
        order = await self.orders.get_by_id(order_id)
        if not order:
            raise HTTPException(status_code=404, detail="Order not found")
        if current_user.id not in {order.customer_id, order.current_executor_id}:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Only order owner or assigned executor can access annotated dataset",
            )
        if order.dataset_link is None:
            raise HTTPException(status_code=400, detail="Order has no dataset")

        dataset = await self.datasets.get_by_id(order.dataset_link)
        if not dataset:
            raise HTTPException(status_code=404, detail="Order dataset not found")

        return order, dataset

    @staticmethod
    def _is_text_order(order: Order) -> bool:
        tags = {tag.lower() for tag in (order.operation_tags or [])}
        return TEXT_DATASET_TAG in tags

    def _validated_objects(self, parts: list[DatasetPart]) -> list[dict[str, str]]:
        objects: list[dict[str, str]] = []
        for part in parts:
            prefix = f"validated/{part.id}/"
            objects.extend(
                {
                    "bucket": part.bucket,
                    "object_key": object_key,
                    "relative_path": object_key.removeprefix(prefix),
                }
                for object_key in self.storage.list_objects(prefix, bucket_name=part.bucket)
            )
        return objects

    def _current_text_rows(self, parts: list[DatasetPart]) -> dict[str, str]:
        rows: dict[str, str] = {}
        for part in parts:
            prefix = f"validated/{part.id}/"
            for object_key in self.storage.list_objects(prefix, bucket_name=part.bucket):
                if object_key.removeprefix(prefix) != "context.csv":
                    continue
                payload = self.storage.download_bytes(
                    object_key,
                    bucket_name=part.bucket,
                ).decode("utf-8-sig", errors="replace")
                reader = csv.DictReader(StringIO(payload))
                if reader.fieldnames != ["text"]:
                    continue
                for index, row in enumerate(reader, start=1):
                    text = (row.get("text") or "").strip()
                    if text:
                        rows[f"{object_key}#row-{index}"] = text
        return rows

    async def _current_labeling_objects(
        self,
        dataset_id: int,
        dataset_objects: list[dict[str, str]],
    ) -> list[str]:
        project = await self.projects.get_by_order_dataset(dataset_id)
        if not project:
            return []

        current_source_keys = {item["object_key"] for item in dataset_objects}
        current_tasks = [
            task
            for task in await self.tasks.list_by_project_id(project.id)
            if task.source_object_key in current_source_keys
        ]
        if not current_tasks:
            return []

        current_task_ids = {task.id for task in current_tasks}
        annotations = await self.annotations.list_by_project_id(project.id)
        annotation_keys = [
            annotation.annotation_object_key
            for annotation in annotations
            if annotation.task_id in current_task_ids and annotation.annotation_object_key
        ]
        return annotation_keys

    async def _list_label_annotation_names(self, order_id: int) -> list[str]:
        base_types = await self.annotation_types.list_by_order(order_id)
        executor_types = await self.executor_annotation_types.list_by_order(order_id)

        result: list[str] = []
        seen: set[str] = set()
        for annotation_type in [*base_types, *executor_types]:
            if annotation_type.annotation in seen:
                continue
            seen.add(annotation_type.annotation)
            result.append(annotation_type.annotation)
        return result

    @staticmethod
    def _unique_archive_path(
        path: PurePosixPath,
        used_paths: set[str],
        used_filenames: dict[str, set[str]],
    ) -> str:
        namespace = path.parts[0] if path.parts else ""
        namespace_filenames = used_filenames.setdefault(namespace, set())
        path_text = path.as_posix()
        if path_text not in used_paths and path.name not in namespace_filenames:
            used_paths.add(path_text)
            namespace_filenames.add(path.name)
            return path_text

        parent = path.parent
        stem = path.stem
        suffix = path.suffix
        counter = 1
        while True:
            candidate_name = f"{stem}_{counter}{suffix}"
            candidate = parent / candidate_name
            candidate_text = candidate.as_posix()
            if candidate_text not in used_paths and candidate_name not in namespace_filenames:
                used_paths.add(candidate_text)
                namespace_filenames.add(candidate_name)
                return candidate_text
            counter += 1

    @staticmethod
    def _extract_preview_annotations(result_payload: dict, annotation_class_ids: dict[str, int]) -> list[dict]:
        results = result_payload.get("result")
        if not isinstance(results, list):
            return []

        preview_annotations: list[dict] = []
        for item in results:
            if not isinstance(item, dict):
                continue

            value = item.get("value")
            if not isinstance(value, dict):
                continue

            labels = value.get("rectanglelabels")
            if not isinstance(labels, list) or not labels:
                continue

            label_name = labels[0]
            if label_name not in annotation_class_ids:
                continue

            x = OrderExportService._normalized_coordinate(value.get("x"))
            y = OrderExportService._normalized_coordinate(value.get("y"))
            width = OrderExportService._normalized_coordinate(value.get("width"))
            height = OrderExportService._normalized_coordinate(value.get("height"))
            if None in {x, y, width, height}:
                continue

            preview_annotations.append(
                {
                    "class_id": annotation_class_ids[label_name],
                    "x_center": round(x + width / 2, 6),
                    "y_center": round(y + height / 2, 6),
                    "width": round(width, 6),
                    "height": round(height, 6),
                }
            )

        return preview_annotations

    @staticmethod
    def _extract_text_sentiment(result_payload: dict) -> str | None:
        results = result_payload.get("result")
        if not isinstance(results, list):
            return None
        for item in results:
            if not isinstance(item, dict):
                continue
            value = item.get("value")
            if not isinstance(value, dict):
                continue
            choices = value.get("choices")
            if isinstance(choices, list) and choices:
                return str(choices[0])
        return None

    @staticmethod
    def _extract_annotator(result_payload: dict, completed_by: str | None) -> int | str | None:
        value = result_payload.get("completed_by")
        if isinstance(value, int):
            return value
        if isinstance(value, dict):
            if value.get("id") is not None:
                return value["id"]
            return value.get("email") or value.get("username")
        if completed_by is not None:
            try:
                return int(completed_by)
            except ValueError:
                return completed_by
        return None

    @staticmethod
    def _payload_timestamp(result_payload: dict, field: str, fallback) -> str:
        value = result_payload.get(field)
        if isinstance(value, str) and value:
            return value
        return fallback.isoformat().replace("+00:00", "Z")

    @staticmethod
    def _normalized_coordinate(value: object) -> float | None:
        if not isinstance(value, (int, float)):
            return None
        return float(value) / 100.0

    @staticmethod
    def _resolve_public_url(configured_url: str | None, request: Request) -> str:
        forwarded_proto = request.headers.get("x-forwarded-proto")
        forwarded_host = request.headers.get("x-forwarded-host")
        scheme = forwarded_proto or request.url.scheme
        host = forwarded_host or request.headers.get("host") or request.url.netloc

        if not configured_url:
            return f"{scheme}://{host}"

        target = configured_url if "://" in configured_url else f"http://{configured_url}"
        parts = urlsplit(target)
        hostname = parts.hostname or ""
        if hostname not in {"localhost", "127.0.0.1", "0.0.0.0"}:
            return configured_url.rstrip("/")

        request_parts = urlsplit(f"{scheme}://{host}")
        resolved_netloc = request_parts.netloc
        if parts.port and request_parts.port != parts.port:
            resolved_host = request_parts.hostname or request_parts.netloc
            resolved_netloc = f"{resolved_host}:{parts.port}"

        return urlunsplit(
            (
                scheme,
                resolved_netloc,
                parts.path.rstrip("/"),
                parts.query,
                parts.fragment,
            )
        ).rstrip("/")