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

backend/src/api/integrations/label_studio.py
186 строк6 KB
AlexanderShmygol
add annotation_executor
03 май 2026, 17:02
03 май 2026, 17:029eaf48f
from dataclasses import dataclass

import httpx

from src.core.config import label_studio_settings


@dataclass
class LabelStudioProjectResult:
    id: int
    raw: dict


@dataclass
class LabelStudioTaskResult:
    id: int
    raw: dict


@dataclass
class LabelStudioWebhookResult:
    id: int
    raw: dict


class LabelStudioGateway:
    def __init__(
        self,
        host: str | None = None,
        username: str | None = None,
        password: str | None = None,
    ) -> None:
        self.host = (host or label_studio_settings.LABEL_STUDIO_API_URL).rstrip("/")
        self.username = username or label_studio_settings.LABEL_STUDIO_USERNAME
        self.password = password or label_studio_settings.LABEL_STUDIO_PASSWORD
        self.client = httpx.Client(
            base_url=self.host,
            timeout=label_studio_settings.LABEL_STUDIO_TIMEOUT_SECONDS,
            follow_redirects=True,
        )
        self._login()

    def create_project(self, title: str, label_config: str) -> LabelStudioProjectResult:
        response = self.client.post(
            "/api/projects/",
            json={
                "title": title,
                "label_config": label_config,
                "show_instruction": False,
                "show_skip_button": True,
                "enable_empty_annotation": True,
            },
            headers=self._api_headers(),
        )
        self._raise_for_status(response, "create project")
        payload = response.json()
        return LabelStudioProjectResult(id=int(payload["id"]), raw=payload)

    def get_project(self, project_id: int) -> LabelStudioProjectResult | None:
        response = self.client.get(
            f"/api/projects/{project_id}",
            headers=self._api_headers(),
        )
        if response.status_code == 404:
            return None
        self._raise_for_status(response, "get project")
        payload = response.json()
        return LabelStudioProjectResult(id=int(payload["id"]), raw=payload)

    def update_project_label_config(
        self,
        project_id: int,
        label_config: str,
    ) -> LabelStudioProjectResult:
        response = self.client.patch(
            f"/api/projects/{project_id}",
            json={"label_config": label_config},
            headers=self._api_headers(),
        )
        self._raise_for_status(response, "update project label config")
        payload = response.json()
        return LabelStudioProjectResult(id=int(payload["id"]), raw=payload)

    def create_task(self, project_id: int, data: dict) -> LabelStudioTaskResult:
        response = self.client.post(
            "/api/tasks/",
            json={"project": project_id, "data": data},
            headers=self._api_headers(),
        )
        self._raise_for_status(response, "create task")
        payload = response.json()
        return LabelStudioTaskResult(id=int(payload["id"]), raw=payload)

    def update_task(self, task_id: int, data: dict) -> LabelStudioTaskResult:
        response = self.client.patch(
            f"/api/tasks/{task_id}",
            json={"data": data},
            headers=self._api_headers(),
        )
        self._raise_for_status(response, "update task")
        payload = response.json()
        return LabelStudioTaskResult(id=int(payload["id"]), raw=payload)

    def delete_task(self, task_id: int) -> None:
        response = self.client.delete(
            f"/api/tasks/{task_id}",
            headers=self._api_headers(),
        )
        if response.status_code == 404:
            return
        self._raise_for_status(response, "delete task")

    def create_webhook(self, project_id: int, url: str) -> LabelStudioWebhookResult:
        response = self.client.post(
            "/api/webhooks/",
            json={
                "url": url,
                "project": project_id,
                "actions": ["ANNOTATION_CREATED", "ANNOTATION_UPDATED", "ANNOTATIONS_DELETED"],
                "send_for_all_actions": False,
                "send_payload": True,
            },
            headers=self._api_headers(),
        )
        self._raise_for_status(response, "create webhook")
        payload = response.json()
        return LabelStudioWebhookResult(id=int(payload["id"]), raw=payload)

    def delete_webhook(self, webhook_id: int) -> None:
        response = self.client.delete(
            f"/api/webhooks/{webhook_id}",
            headers=self._api_headers(),
        )
        if response.status_code == 404:
            return
        self._raise_for_status(response, "delete webhook")

    def close(self) -> None:
        self.client.close()

    def _login(self) -> None:
        login_page = self.client.get("/user/login/")
        self._raise_for_status(login_page, "load login page")

        csrf_token = self.client.cookies.get("csrftoken")
        if not csrf_token:
            raise RuntimeError("Label Studio login page did not provide CSRF cookie")

        response = self.client.post(
            "/user/login/?next=/",
            data={
                "csrfmiddlewaretoken": csrf_token,
                "email": self.username,
                "password": self.password,
                "persist_session": "on",
            },
            headers={
                "Referer": f"{self.host}/user/login/",
                "X-CSRFToken": csrf_token,
            },
        )
        self._raise_for_status(response, "login to Label Studio")

        whoami = self.client.get("/api/projects/", headers=self._api_headers())
        self._raise_for_status(whoami, "verify Label Studio session")

    def _api_headers(self) -> dict[str, str]:
        csrf_token = self.client.cookies.get("csrftoken", "")
        headers = {
            "Referer": f"{self.host}/",
            "Accept": "application/json",
        }
        if csrf_token:
            headers["X-CSRFToken"] = csrf_token
        return headers

    @staticmethod
    def _raise_for_status(response: httpx.Response, action: str) -> None:
        if response.is_success:
            return
        details = response.text
        try:
            details = str(response.json())
        except Exception:
            pass
        raise RuntimeError(f"Failed to {action} in Label Studio: {response.status_code} {details}")