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

backend/src/api/services/label_studio.py
373 строки16 KB
AlexanderShmygol
fixed
17 май 2026, 19:27
17 май 2026, 19:2749e260c
import csv
from io import StringIO
from xml.sax.saxutils import escape
from urllib.parse import urlsplit, urlunsplit

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

from src.api.integrations.label_studio import LabelStudioGateway
from src.api.repositories.label_studio_instance import LabelStudioInstanceRepository
from src.api.repositories.label_studio_project import LabelStudioProjectRepository
from src.api.repositories.order_annotation_type import OrderAnnotationTypeRepository
from src.api.repositories.order_annotation_type_executor import OrderAnnotationTypeExecutorRepository
from src.api.repositories.label_studio_task import LabelStudioTaskRepository
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 label_studio_settings, minio_settings
from src.models.label_studio_instance import LabelStudioInstance
from src.models.label_studio_project import LabelStudioProject
from src.models.label_studio_task import LabelStudioTask
from src.models.order_annotation_type import OrderAnnotationType, OrderAnnotationTypeExecutor
from src.models.order_dataset import DatasetPart, OrderDatasetValidationStatus
from src.models.user import User

TEXT_DATASET_TAG = "text"
IMAGE_DATASET_TAG = "images"


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

    @staticmethod
    def _has_workspace_access(order_customer_id: int, order_executor_id: int | None, user_id: int) -> bool:
        return user_id in {order_customer_id, order_executor_id}

    async def open_workspace(self, order_id: int, current_user: User, request: Request) -> LabelStudioProject:
        order = await self.orders.get_by_id(order_id)
        if not order:
            raise HTTPException(status_code=404, detail="Order not found")
        if not self._has_workspace_access(order.customer_id, order.current_executor_id, current_user.id):
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Only order owner or assigned executor can open labeling workspace",
            )
        if order.dataset_link is None:
            raise HTTPException(status_code=400, detail="Order has no uploaded dataset")

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

        current_parts = await self.datasets.list_current_parts(dataset.id)
        if not current_parts:
            raise HTTPException(status_code=400, detail="Dataset has no parts")
        if not all(part.validation_status == OrderDatasetValidationStatus.VALIDATION for part in current_parts):
            raise HTTPException(
                status_code=400,
                detail="Dataset must be validated before labeling",
            )
        annotation_types = await self._list_label_annotation_types(order.id)
        if not annotation_types:
            raise HTTPException(status_code=400, detail="Order annotation types are not configured")
        dataset_kind = self._dataset_kind(order.operation_tags)
        label_config = self._build_label_config(annotation_types, dataset_kind)
        public_label_studio_url = self._resolve_public_url(
            label_studio_settings.LABEL_STUDIO_PUBLIC_URL,
            request,
        )
        public_minio_endpoint = self._resolve_public_url(
            minio_settings.MINIO_PUBLIC_ENDPOINT,
            request,
        )

        existing_project = await self.projects.get_by_order_dataset(dataset.id)

        instance = await self._get_or_create_instance()
        gateway = LabelStudioGateway(
            host=instance.host,
            username=instance.username,
            password=instance.password,
        )

        project_title = f"Order {order.id} Dataset {dataset.id}"
        webhook_url = f"{label_studio_settings.LABEL_STUDIO_WEBHOOK_BASE_URL.rstrip('/')}/events"
        try:
            remote_label_config: str | None = None
            if existing_project is not None:
                remote_project = gateway.get_project(existing_project.external_project_id)
                if remote_project is None:
                    await self.projects.delete(existing_project)
                    existing_project = None
                else:
                    remote_label_config = remote_project.raw.get("label_config")

            if existing_project is None:
                created_project = gateway.create_project(
                    title=project_title,
                    label_config=label_config,
                )
                workspace_url = (
                    f"{public_label_studio_url.rstrip('/')}/projects/{created_project.id}/data/"
                )
                export_prefix = f"orders/{order.id}/datasets/{dataset.id}/labeling"
                project = await self.projects.create(
                    LabelStudioProject(
                        instance_id=instance.id,
                        order_id=order.id,
                        order_dataset_id=dataset.id,
                        external_project_id=created_project.id,
                        title=project_title,
                        workspace_url=workspace_url,
                        label_config=label_config,
                        export_prefix=export_prefix,
                    )
                )
            else:
                project = existing_project
                project.workspace_url = (
                    f"{public_label_studio_url.rstrip('/')}/projects/{project.external_project_id}/data/"
                )
                if project.label_config != label_config or (
                    remote_label_config is not None and remote_label_config != label_config
                ):
                    gateway.update_project_label_config(project.external_project_id, label_config)
                    project.label_config = label_config
                project = await self.projects.save(project)

            if project.external_webhook_id is not None:
                try:
                    gateway.delete_webhook(project.external_webhook_id)
                except RuntimeError:
                    pass

            webhook = gateway.create_webhook(
                project_id=project.external_project_id,
                url=webhook_url,
            )
            project.external_webhook_id = webhook.id
            project = await self.projects.save(project)

            existing_tasks = {
                task.source_object_key: task
                for task in await self.tasks.list_by_project_id(project.id)
            }
            validated_objects = self._validated_objects(current_parts, dataset_kind)
            current_source_keys = {item["object_key"] for item in validated_objects}

            for object_key, task in list(existing_tasks.items()):
                if object_key in current_source_keys:
                    continue
                gateway.delete_task(task.external_task_id)
                await self.tasks.delete(task)
                existing_tasks.pop(object_key, None)

            for item in validated_objects:
                object_key = item["object_key"]
                task_data = self._task_data(item, dataset_kind, public_minio_endpoint)
                source_url = item.get("source_url", "")
                if object_key in existing_tasks:
                    task = existing_tasks[object_key]
                    if task.source_url != source_url:
                        gateway.update_task(task.external_task_id, task_data)
                        task.source_url = source_url
                        await self.tasks.save(task)
                    continue
                created_task = gateway.create_task(
                    project_id=project.external_project_id,
                    data=task_data,
                )
                await self.tasks.create(
                    LabelStudioTask(
                        project_id=project.id,
                        external_task_id=created_task.id,
                        source_object_key=object_key,
                        source_url=source_url,
                    )
                )
        except RuntimeError as exc:
            raise HTTPException(
                status_code=status.HTTP_502_BAD_GATEWAY,
                detail=str(exc),
            ) from exc
        finally:
            gateway.close()

        return project

    async def count_project_tasks(self, project_id: int) -> int:
        return await self.tasks.count_by_project_id(project_id)

    async def _list_label_annotation_types(
        self,
        order_id: int,
    ) -> list[OrderAnnotationType | OrderAnnotationTypeExecutor]:
        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[OrderAnnotationType | OrderAnnotationTypeExecutor] = []
        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)
        return result

    async def _get_or_create_instance(self) -> LabelStudioInstance:
        instance = await self.instances.get_by_host(label_studio_settings.LABEL_STUDIO_API_URL.rstrip("/"))
        if instance:
            instance.username = label_studio_settings.LABEL_STUDIO_USERNAME
            instance.password = label_studio_settings.LABEL_STUDIO_PASSWORD
            instance.user_token = label_studio_settings.LABEL_STUDIO_USER_TOKEN
            instance.is_active = True
            return await self.instances.save(instance)

        return await self.instances.create(
            LabelStudioInstance(
                host=label_studio_settings.LABEL_STUDIO_API_URL.rstrip("/"),
                username=label_studio_settings.LABEL_STUDIO_USERNAME,
                password=label_studio_settings.LABEL_STUDIO_PASSWORD,
                user_token=label_studio_settings.LABEL_STUDIO_USER_TOKEN,
                is_active=True,
            )
        )

    @staticmethod
    def _dataset_kind(operation_tags: list[str] | None) -> str:
        tags = {tag.lower() for tag in (operation_tags or [])}
        if TEXT_DATASET_TAG in tags:
            return TEXT_DATASET_TAG
        return IMAGE_DATASET_TAG

    def _validated_objects(self, parts: list[DatasetPart], dataset_kind: str) -> list[dict[str, str]]:
        if dataset_kind == TEXT_DATASET_TAG:
            return self._validated_text_objects(parts)

        objects: list[dict[str, str]] = []
        for part in parts:
            prefix = f"validated/{part.id}/"
            objects.extend(
                {"bucket": part.bucket, "object_key": object_key}
                for object_key in self.storage.list_objects(prefix, bucket_name=part.bucket)
            )
        if not objects:
            raise HTTPException(status_code=400, detail="Validated dataset assets not found")
        return objects

    def _validated_text_objects(self, parts: list[DatasetPart]) -> list[dict[str, str]]:
        contexts: list[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):
                relative_key = object_key.removeprefix(prefix)
                if relative_key == "context.csv":
                    contexts.extend(self._text_context_rows(part.bucket, object_key))

        if not contexts:
            raise HTTPException(status_code=400, detail="Validated text contexts not found")

        return contexts

    def _text_context_rows(self, bucket: str, object_key: str) -> list[dict[str, str]]:
        payload = self.storage.download_bytes(object_key, bucket_name=bucket).decode("utf-8-sig", errors="replace")
        reader = csv.DictReader(StringIO(payload))
        if reader.fieldnames != ["text"]:
            raise HTTPException(status_code=400, detail="Validated context.csv has invalid columns")

        rows: list[dict[str, str]] = []
        for index, row in enumerate(reader, start=1):
            context_text = (row.get("text") or "").strip()
            if not context_text:
                continue
            rows.append(
                {
                    "bucket": bucket,
                    "object_key": f"{object_key}#row-{index}",
                    "context_text": context_text,
                }
            )
        return rows

    def _task_data(
        self,
        item: dict[str, str],
        dataset_kind: str,
        public_minio_endpoint: str,
    ) -> dict[str, str]:
        if dataset_kind == TEXT_DATASET_TAG:
            return {
                "text": item["context_text"],
            }
        source_url = self.storage.get_presigned_get_url(
            item["object_key"],
            bucket_name=item["bucket"],
            public_endpoint=public_minio_endpoint,
        )
        item["source_url"] = source_url
        return {"image": source_url}

    @staticmethod
    def _build_label_config(
        annotation_types: list[OrderAnnotationType | OrderAnnotationTypeExecutor],
        dataset_kind: str,
    ) -> str:
        labels = "\n".join(
            f'    <Label value="{escape(annotation_type.annotation)}"/>'
            for annotation_type in annotation_types
        )
        if dataset_kind == TEXT_DATASET_TAG:
            choices = "\n".join(
                f'    <Choice value="{escape(annotation_type.annotation)}"/>'
                for annotation_type in annotation_types
            )
            return (
                "<View>\n"
                '  <Header value="Situation"/>\n'
                '  <Text name="text" value="$text"/>\n'
                '  <Choices name="label" toName="text" choice="single">\n'
                f"{choices}\n"
                "  </Choices>\n"
                "</View>"
            )
        return (
            "<View>\n"
            '  <Image name="image" value="$image"/>\n'
            '  <RectangleLabels name="label" toName="image">\n'
            f"{labels}\n"
            "  </RectangleLabels>\n"
            "</View>"
        )

    @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("/")