intellif-aihub 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

aihub/__init__.py

@@ -1,4 +1 @@
-from .stop import is_pre_stopped, on_pre_stop
-from ._version import __version__
-
-__all__ = ["is_pre_stopped", "on_pre_stop", "__version__"]
+__version__ = "0.1.2"

aihub/client.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import os
+
+import httpx
+
+from .exceptions import APIError
+from .services import artifact
+from .services import dataset_management
+from .services import document_center
+from .services import labelfree
+from .services import tag_management
+from .services import task_center
+from .services.artifact import ArtifactService
+from .services.dataset_management import DatasetManagementService
+from .services.document_center import DocumentCenterService
+from .services.labelfree import LabelfreeService
+from .services.tag_management import TagManagementService
+from .services.task_center import TaskCenterService
+
+
+class Client:
+    """AI-HUB python SDK 客户端
+
+    Attributes:
+        dataset_management (DatasetManagementService): 数据集管理服务
+        labelfree (LabelfreeService): 标注服务
+        task_center (TaskCenterService): 任务中心
+        artifact (ArtifactService): 制品管理
+
+    """
+
+    labelfree: LabelfreeService = None
+    tag_management: TagManagementService = None
+    document_center: DocumentCenterService = None
+    task_center: TaskCenterService = None
+    dataset_management: DatasetManagementService = None
+    artifact: ArtifactService = None
+
+    def __init__(
+        self, *, base_url: str, token: str | None = None, timeout: float = 60.0
+    ):
+        """AI-HUB python SDK 客户端
+
+        Args:
+            base_url (str): 服务地址
+            token (str): 密钥，显式传入，或在环境变量AI_HUB_TOKEN中设置
+
+        Examples:
+            >>> from aihub.client import Client
+            >>> client = Client(base_url="xxx", token="xxxx")
+
+        """
+        if not base_url:
+            raise ValueError("base_url必须填写")
+
+        token = token or os.getenv("AI_HUB_TOKEN")
+        if not token:
+            raise ValueError("缺少token：请显式传入，或在环境变量AI_HUB_TOKEN中设置")
+
+        self._http = httpx.Client(
+            base_url=base_url.rstrip("/"),
+            timeout=timeout,
+            headers={"Authorization": f"Bearer {token}"},
+            # event_hooks={"response": [self._raise_for_status]},
+        )
+        self.dataset_management = dataset_management.DatasetManagementService(
+            self._http
+        )
+        self.labelfree = labelfree.LabelfreeService(self._http)
+        self.tag_management = tag_management.TagManagementService(self._http)
+        self.document_center = document_center.DocumentCenterService(self._http)
+        self.task_center = task_center.TaskCenterService(self._http)
+        self.artifact = artifact.ArtifactService(self._http)
+
+    @staticmethod
+    def _raise_for_status(r: httpx.Response):
+        try:
+            r.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise APIError(f"{e.response.status_code}: {e.response.text}") from e
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self._http.close()

aihub/exceptions.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class SDKError(Exception):
+    pass
+
+
+class APIError(SDKError):
+    def __init__(self, message: str, *, status: Optional[int] = None, detail: Any = None) -> None:
+        super().__init__(message)
+        self.message: str = message
+        self.status: Optional[int] = status
+        self.detail: Any = detail
+
+    def __str__(self) -> str:
+        return f"[HTTP {self.status}] {self.message}" if self.status else self.message

aihub/models/artifact.py

@@ -0,0 +1,137 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+"""制品管理模型模块
+
+该模块定义了制品管理相关的数据模型，包括制品类型、创建制品请求、制品响应等。
+"""
+from __future__ import annotations
+
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class ArtifactType(str, Enum):
+    """制品类型枚举
+
+    定义了系统支持的各种制品类型。
+    """
+
+    dataset = "dataset"  # 数据集类型
+    model = "model"  # 模型类型
+    metrics = "metrics"  # 指标类型
+    log = "log"  # 日志类型
+    checkpoint = "checkpoint"  # 检查点类型
+    image = "image"  # 图像类型
+    prediction = "prediction"  # 预测结果类型
+    other = "other"  # 其他类型
+
+
+class CreateArtifactsReq(BaseModel):
+    """创建制品请求模型
+
+    用于向服务器发送创建制品的请求。
+    """
+
+    entity_id: str
+    """实体ID，通常是运行ID，用于关联制品与特定运行"""
+
+    entity_type: ArtifactType = ArtifactType.other
+    """制品类型，指定制品的类型，默认为other"""
+
+    src_path: str
+    """源路径，制品在系统中的路径标识"""
+
+    is_dir: bool = False
+    """是否为目录，True表示制品是一个目录，False表示是单个文件"""
+
+
+class CreateArtifactsResponseData(BaseModel):
+    """创建制品响应数据模型
+
+    服务器创建制品后返回的数据。
+    """
+
+    id: int  # 制品ID
+    s3_path: str  # S3存储路径
+
+
+class CreateArtifactsResponseModel(BaseModel):
+    """创建制品响应模型
+
+    服务器对创建制品请求的完整响应。
+    """
+
+    code: int  # 响应码，0表示成功
+    msg: str = ""  # 响应消息
+    data: CreateArtifactsResponseData | None = None  # 响应数据
+
+
+class CreatEvalReq(BaseModel):
+    """创建评估请求模型
+
+    用于创建模型评估的请求。
+    """
+
+    dataset_id: int  # 数据集ID
+    dataset_version_id: int  # 数据集版本ID
+    prediction_artifact_path: str  # 预测结果制品路径
+    evaled_artifact_path: str  # 评估结果制品路径
+    run_id: str  # 运行ID
+    user_id: int  # 用户ID
+    report: dict = {}  # 评估报告
+
+
+class ArtifactResp(BaseModel):
+    """制品响应模型
+
+    表示一个制品的详细信息。
+    """
+
+    id: int  # 制品ID
+    entity_type: str  # 实体类型
+    entity_id: str  # 实体ID
+    src_path: str  # 源路径
+    s3_path: str  # S3存储路径
+    is_dir: bool  # 是否为目录
+
+
+class ArtifactRespData(BaseModel):
+    """制品响应数据模型
+
+    包含分页信息和制品列表的响应数据。
+    """
+
+    total: int  # 总记录数
+    page_size: int  # 每页大小
+    page_num: int  # 页码
+    data: List[ArtifactResp]  # 制品列表
+
+
+class ArtifactRespModel(BaseModel):
+    """制品响应模型
+
+    服务器对获取制品请求的完整响应。
+    """
+
+    code: int  # 响应码，0表示成功
+    msg: str = ""  # 响应消息
+    data: ArtifactRespData  # 响应数据
+
+
+# 无限大的页面大小，用于一次性获取所有制品
+InfinityPageSize = 10000 * 100
+
+
+class StsResp(BaseModel):
+    """STS临时凭证响应模型
+
+    包含访问S3存储所需的临时凭证信息。
+    """
+
+    access_key_id: Optional[str] = None  # 访问密钥ID
+    secret_access_key: Optional[str] = None  # 秘密访问密钥
+    session_token: Optional[str] = None  # 会话令牌
+    expiration: Optional[int] = None  # 过期时间
+    endpoint: Optional[str] = None  # 端点URL

aihub/models/common.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Generic, Optional, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T")
+
+
+class APIWrapper(BaseModel, Generic[T]):
+    code: int
+    msg: Optional[str] = None
+    data: Optional[T]

aihub/models/dataset_management.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class CreateDatasetRequest(BaseModel):
+    name: str
+    description: str
+    tags: List[int]
+    cover_img: Optional[str] = Field(None, alias="cover_img")
+    create_by: Optional[int] = Field(None, alias="create_by")
+    is_private: Optional[bool] = Field(None, alias="is_private")
+    access_user_ids: Optional[List[int]] = Field(None, alias="access_user_ids")
+
+
+class CreateDatasetResponse(BaseModel):
+    id: int = Field(alias="id")
+
+
+class DatasetVersionBase(BaseModel):
+    id: int
+    version: int
+    status: int
+    parquet_index_path: Optional[str] = Field(None, alias="parquet_index_path")
+    data_count: int = Field(alias="data_count")
+
+
+class DatasetDetail(BaseModel):
+    id: int
+    name: str
+    description: str
+    cover_img: Optional[str] = Field(None, alias="cover_img")
+    created_at: int = Field(alias="created_at")
+    updated_at: int = Field(alias="update_at")
+    user_id: int = Field(alias="user_id")
+    username: str
+    tags: List[int]
+    access_user_ids: Optional[List[int]] = Field(None, alias="access_user_ids")
+    is_private: Optional[bool] = Field(None, alias="is_private")
+    versions: List[DatasetVersionBase]
+
+
+class ExtInfo(BaseModel):
+    rec_file_path: Optional[str] = Field(None, alias="rec_file_path")
+    idx_file_path: Optional[str] = Field(None, alias="idx_file_path")
+    json_file_path: Optional[str] = Field(None, alias="json_file_path")
+    image_dir_path: Optional[str] = Field(None, alias="image_dir_path")
+
+
+class CreateDatasetVersionRequest(BaseModel):
+    upload_path: str = Field(alias="upload_path")
+    description: Optional[str] = None
+    dataset_id: int = Field(alias="dataset_id")
+    object_cnt: Optional[int] = Field(None, alias="object_cnt")
+    data_size: Optional[int] = Field(None, alias="data_size")
+    create_by: Optional[int] = Field(None, alias="create_by")
+    upload_type: Optional[int] = Field(4, alias="upload_type")
+    ext_info: Optional[ExtInfo] = Field(None, alias="ext_info")
+
+
+class CreateDatasetVersionResponse(BaseModel):
+    id: int = Field(alias="id")
+
+
+class UploadDatasetVersionRequest(BaseModel):
+    upload_path: str = Field(alias="upload_path")
+    upload_type: int = Field(alias="upload_type")
+    dataset_id: int = Field(alias="dataset_id")
+    parent_version_id: Optional[int] = Field(None, alias="parent_version_id")
+    description: Optional[str] = Field(None, alias="description")
+
+
+class UploadDatasetVersionResponse(BaseModel):
+    id: int = Field(alias="id")
+
+
+class DatasetVersionDetail(BaseModel):
+    id: int
+    version: int
+    dataset_id: int = Field(alias="dataset_id")
+    upload_path: str = Field(alias="upload_path")
+    upload_type: int = Field(alias="upload_type")
+    parent_version_id: Optional[int] = Field(None, alias="parent_version_id")
+    description: Optional[str] = None
+    status: int
+    message: Optional[str] = None
+    created_at: int = Field(alias="created_at")
+    user_id: int = Field(alias="user_id")
+    data_size: Optional[int] = Field(None, alias="data_size")
+    data_count: Optional[int] = Field(None, alias="data_count")
+    parquet_index_path: Optional[str] = Field(None, alias="parquet_index_path")
+    ext_info: Optional[ExtInfo] = Field(None, alias="ext_info")
+
+
+class FileUploadData(BaseModel):
+    path: str
+    url: str

aihub/models/document_center.py

@@ -0,0 +1,28 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+from __future__ import annotations
+
+from typing import List, Optional, Any
+
+from pydantic import BaseModel
+
+
+class Document(BaseModel):
+    id: int
+    title: str
+    type: int
+    edit_time: int
+    need_update: bool
+    content: str
+    username: str
+    user_id: int
+    created_at: int
+    comments: Optional[Any] = None
+
+
+class GetDocumentsResponse(BaseModel):
+    total: int
+    page_size: int
+    page_num: int
+    data: List[Document]

aihub/models/labelfree.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class Stats(BaseModel):
+    """标注统计信息"""
+
+    total_annotations: int = Field(alias="total_annotations")
+    labeled_annotations: int = Field(alias="labeled_annotations")
+    total_labels: int = Field(alias="total_labels")
+    total_reviews: Optional[int] = Field(None, alias="total_reviews")
+    unlabeled_reviews: Optional[int] = Field(None, alias="unlabeled_reviews")
+    labeled_reviews: Optional[int] = Field(None, alias="labeled_reviews")
+    accepted_count: Optional[int] = Field(None, alias="accepted_count")
+    rejected_count: Optional[int] = Field(None, alias="rejected_count")
+
+
+class GetGlobalStatsResponse(BaseModel):
+    """
+    标注统计概况
+    """
+
+    global_stats: Stats = Field(alias="global_stats")
+    valid_ten_percent: bool = Field(alias="valid_ten_percent")
+    valid_fifty_percent: bool = Field(alias="valid_fifty_percent")
+    valid_hundred_percent: bool = Field(alias="valid_hundred_percent")
+    data_exported_count: int = Field(alias="data_exported_count")
+    exported_dataset_name: str = Field(alias="exported_dataset_name")

aihub/models/tag_management.py

@@ -0,0 +1,21 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+from __future__ import annotations
+
+from typing import List
+
+from pydantic import BaseModel
+
+
+class Project(BaseModel):
+    id: int
+    name: str
+
+
+class ProjectListData(BaseModel):
+    data: List[Project]
+
+
+class SelectProjectsResponse(BaseModel):
+    data: List[Project]

aihub/models/task_center.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import json
+from enum import Enum
+from typing import Optional, List
+
+from pydantic import BaseModel, Field, field_serializer, field_validator
+
+
+class TaskCenterPriorityEnum(Enum):
+    """
+    任务优先级枚举
+    """
+
+    low = "low"
+    medium = "medium"
+    high = "high"
+
+
+class LabelProjectTypeEnum(Enum):
+    """
+    任务类型枚举
+    """
+
+    IMAGE_CLASSIFICATION = 1
+    OBJECT_DETECTION = 2
+    SEMANTIC_SEGMENTATION = 3
+    TEXT_CLASSIFICATION = 4
+    TEXT_SEGMENTATION = 5
+    TEXT_RECOGNITION = 6
+    IMAGE_CAPTIONING = 7
+    IMAGE_STYLE_TRANSFER = 8
+    IMAGE_COLORIZATION = 9
+    IMAGE_RESTORATION = 10
+
+
+class CreateTaskOtherInfo(BaseModel):
+    label_project_type: LabelProjectTypeEnum = LabelProjectTypeEnum.IMAGE_CLASSIFICATION
+    dataset_id: int = Field(alias="dataset_id")
+    dataset_version_id: int = Field(alias="dataset_version_id")
+    doc_id: int = Field(alias="doc_id")
+    doc_type: str = Field(alias="doc_type", default="doc_center")
+
+
+class ProjectInfo(BaseModel):
+    label_project_id: int = Field(alias="label_project_id")
+    label_project_name: str = Field(alias="label_project_name")
+
+
+class TaskDetailOtherInfo(BaseModel):
+    label_project_type: LabelProjectTypeEnum = LabelProjectTypeEnum.IMAGE_CLASSIFICATION
+    dataset_id: int = Field(alias="dataset_id")
+    dataset_version_id: int = Field(alias="dataset_version_id")
+    doc_id: int = Field(alias="doc_id")
+    doc_type: str = Field(alias="doc_type", default="doc_center")
+    label_projects: Optional[List[ProjectInfo]] = Field(alias="label_projects")
+
+
+class CreateTaskReq(BaseModel):
+    name: str
+    description: Optional[str] = None
+    task_priority: Optional[str] = None
+    type: Optional[str] = None
+    receiver_id: Optional[int] = None
+    project_id: Optional[int] = None
+    other_info: CreateTaskOtherInfo = Field(alias="other_info")
+    estimated_delivery_at: Optional[int] = None
+
+    @field_serializer("other_info")
+    def serialize_other_info(self, value: CreateTaskOtherInfo) -> str:
+        """将 other_info 序列化为 JSON 字符串"""
+        return value.model_dump_json()
+
+
+class CreateTaskResp(BaseModel):
+    id: int = Field(alias="id")
+
+
+class LabelTaskDetail(BaseModel):
+    """任务详情"""
+
+    name: str
+    description: Optional[str] = Field(alias="description")
+    task_priority: Optional[str] = Field(alias="task_priority")
+    type: Optional[str] = Field(alias="type")
+    receiver_id: Optional[int] = Field(alias="receiver_id")
+    project_id: Optional[int] = None
+    other_info: TaskDetailOtherInfo = Field(alias="other_info")
+    estimated_delivery_at: Optional[int] = None
+
+    @field_serializer("other_info")
+    def serialize_other_info(self, value: TaskDetailOtherInfo) -> str:
+        """将 other_info 序列化为 JSON 字符串"""
+        return value.model_dump_json()
+
+    @field_validator("other_info", mode="before")
+    @classmethod
+    def parse_other_info(cls, value):
+        """将字符串解析为 TaskDetailOtherInfo 对象"""
+        if isinstance(value, str):
+            try:
+                # 解析 JSON 字符串为字典
+                data = json.loads(value)
+                # 创建 TaskDetailOtherInfo 对象
+                return TaskDetailOtherInfo(**data)
+            except (json.JSONDecodeError, TypeError, ValueError) as e:
+                raise ValueError(f"无法解析 other_info 字符串: {e}")
+        elif isinstance(value, dict):
+            # 如果传入的是字典，直接创建对象
+            return TaskDetailOtherInfo(**value)
+        elif isinstance(value, TaskDetailOtherInfo):
+            # 如果已经是对象，直接返回
+            return value
+        else:
+            raise ValueError(
+                f"other_info 必须是字符串、字典或 TaskDetailOtherInfo 对象，得到: {type(value)}"
+            )

aihub/models/user.py

@@ -0,0 +1,46 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class UserSearchReq(BaseModel):
+    page_size: int = 1000000
+    page_num: int = 1
+    username: str = ""
+    nickname: str = ""
+    email: str = ""
+    user_ids: List[int] = []
+    role_ids: List[int] = []
+    role_names: List[str] = []
+    status: int = 1
+
+
+class Role(BaseModel):
+    id: Optional[int] = None
+    name: Optional[str] = None
+    role_type: Optional[int] = None
+    menu_ids: Optional[List[int]] = None
+
+
+class UserSearchDatum(BaseModel):
+    id: Optional[int] = 0
+    username: Optional[str] = None
+    nickname: Optional[str] = None
+    email: Optional[str] = None
+    status: Optional[int] = None
+    created_at: Optional[int] = None
+    roles: Optional[List[Role]] = None
+    tags: Optional[List[int]] = None
+
+
+class UserSearchListData(BaseModel):
+    total: int
+    page_size: int
+    page_num: int
+    data: List[UserSearchDatum]