from collections.abc import Generator
from datetime import datetime
from dataclasses import asdict
from typing import Any
from chariot import _apis, mcp_setting
from chariot.datasets import _utils, models
from chariot_api._openapi.datasets_v3 import models as openapi_models
__all__ = [
"create_annotation",
"get_annotation",
"get_datum_annotations",
"get_annotation_history",
"update_annotation",
"archive_annotation",
]
[docs]
@mcp_setting(mutating=True)
def create_annotation(
datum_id: str,
*,
class_label: str | None = None,
contour: list[list[models.Point]] | None = None,
bbox: models.BoundingBox | None = None,
oriented_bbox: models.OrientedBoundingBox | None = None,
text_classification: models.TextClassification | None = None,
text_generation: models.TextGenerationInput | None = None,
token_classification: models.TokenClassification | None = None,
metadata: dict[str, Any] | None = None,
approval_status: models.ApprovalStatus | None = None,
task_id: str | None = None,
) -> models.Annotation:
"""Create a new annotation
:param datum_id: Id of datum to add annotation to
:type datum_id: str
:param class_label: Class label of the annotation
:type class_label: Optional[str]
:param contour: Contour for an Image Segmentation annotation
:type contour: Optional[List[List[models.Point]]]
:param bbox: Bounding box for an Object Detection annotation
:type bbox: Optional[models.BoundingBox]
:param oriented_bbox: Oriented bounding box for an Oriented Object Detection annotation
:type oriented_bbox: Optional[models.OrientedBoundingBox]
:param text_classification: Text Classification annotation
:type text_classification: Optional[models.TextClassification]
:param text_generation: Text Generation annotation
:type text_generation: Optional[models.TextGenerationInput]
:param token_classification: Token Classification annotation
:type token_classification: Optional[models.TokenClassification]
:param metadata: Metadata associated with the annotation
:type metadata: Optional[Dict[str, Any]]
:param approval_status: Reviewer approval status for the annotation
:type approval_status: Optional[models.ApprovalStatus]
:param task_id: Id of task to which the datum is locked with
:type task_id: Optional[str]
:return: New annotation details
:rtype: models.Annotation
"""
request = openapi_models.InputCreateAnnotationRequest(
class_label=class_label,
bbox=_utils.convert_from_dataclass(bbox, openapi_models.ModelBoundingBox),
contour=[_utils.convert_from_dataclass_list(c, openapi_models.ModelPoint) for c in contour]
if contour is not None
else None,
oriented_bbox=_utils.convert_from_dataclass(
oriented_bbox, openapi_models.ModelOrientedBoundingBox
),
text_classification=_utils.convert_from_dataclass(
text_classification, openapi_models.ModelTextClassification
),
text_generation=_utils.convert_from_dataclass(
text_generation, openapi_models.InputTextGenerationInput
),
token_classification=_utils.convert_from_dataclass(
token_classification, openapi_models.ModelTokenClassification
),
metadata=metadata,
approval_status=_utils.enum_value(approval_status),
task_id=task_id,
)
response = _apis.datasets_v3.annotations_api.create_annotation(datum_id=datum_id, body=request)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from create_annotation")
return _utils.convert_to_dataclass(response.data.model_dump(), models.Annotation)
[docs]
def get_annotation(id: str) -> models.Annotation:
"""Get an annotation by id
:param id: Id of annotation to get
:type id: str
:return: Annotation details
:rtype: models.Annotation
"""
response = _apis.datasets_v3.annotations_api.get_annotation(annotation_id=id)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from get_annotation")
return _utils.convert_to_dataclass(response.data.model_dump(), models.Annotation)
[docs]
def get_annotation_history(id: str) -> models.AnnotationHistoryDetails:
"""Get an annotation's history by id.
:param id: Id of annotation to get history for
:type id: str
:return: Annotation history details
:rtype: models.AnnotationHistoryDetails
"""
response = _apis.datasets_v3.annotations_api.get_annotation_history(annotation_id=id)
if not response.data:
raise RuntimeError(
"Received malformed response (missing `data`) from get_annotation_history"
)
return _utils.convert_to_dataclass(response.data.model_dump(), models.AnnotationHistoryDetails)
[docs]
@mcp_setting(mutating=True)
def update_annotation(
annotation_id: str,
*,
class_label: str | None = None,
contour: list[list[models.Point]] | None = None,
bbox: models.BoundingBox | None = None,
oriented_bbox: models.OrientedBoundingBox | None = None,
text_classification: models.TextClassification | None = None,
text_generation: models.TextGenerationInput | None = None,
token_classification: models.TokenClassification | None = None,
metadata: dict[str, Any] | None = None,
approval_status: models.ApprovalStatus | None = None,
updated_at: str | None = None,
task_id: str | None = None,
) -> models.Annotation:
"""Update or replace an annotation
:param annotation_id: Id of annotation to be updted or replaced
:type annotation_id: str
:param class_label: Class label of the annotation
:type class_label: Optional[str]
:param contour: Contour for an Image Segmentation annotation
:type contour: Optional[List[List[models.Point]]]
:param bbox: Bounding box for an Object Detection annotation
:type bbox: Optional[models.BoundingBox]
:param oriented_bbox: Oriented bounding box for an Oriented Object Detection annotation
:type oriented_bbox: Optional[models.OrientedBoundingBox]
:param text_classification: Text Classification annotation
:type text_classification: Optional[models.TextClassification]
:param text_generation: Text Generation annotation
:type text_generation: Optional[models.TextGenerationInput]
:param token_classification: Token Classification annotation
:type token_classification: Optional[models.TokenClassification]
:param metadata: Metadata associated with the annotation
:type metadata: Optional[Dict[str, Any]]
:param approval_status: Reviewer approval status for the annotation
:type approval_status: Optional[models.ApprovalStatus]
:param updated_at: must match the updated_at time on the annotation being updated
:type updated_at: Optional[str]
:param task_id: Id of task to which the datum is locked with
:type task_id: Optional[str]
:return: New annotation details
:rtype: models.Annotation
"""
request = openapi_models.InputUpdateAnnotationRequest(
class_label=class_label,
bbox=_utils.convert_from_dataclass(bbox, openapi_models.ModelBoundingBox),
contour=[_utils.convert_from_dataclass_list(c, openapi_models.ModelPoint) for c in contour]
if contour is not None
else None,
oriented_bbox=_utils.convert_from_dataclass(
oriented_bbox, openapi_models.ModelOrientedBoundingBox
),
text_classification=_utils.convert_from_dataclass(
text_classification, openapi_models.ModelTextClassification
),
text_generation=_utils.convert_from_dataclass(
text_generation, openapi_models.InputTextGenerationInput
),
token_classification=_utils.convert_from_dataclass(
token_classification, openapi_models.ModelTokenClassification
),
metadata=metadata,
approval_status=_utils.enum_value(approval_status),
updated_at=updated_at,
task_id=task_id,
)
response = _apis.datasets_v3.annotations_api.update_annotation(
annotation_id=annotation_id, body=request
)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from update_annotation")
return _utils.convert_to_dataclass(response.data.model_dump(), models.Annotation)
[docs]
@mcp_setting(mutating=True)
def archive_annotation(
id: str,
*,
task_id: str | None = None,
) -> models.Annotation:
"""Archive (soft-delete) an annotation by id
:param id: Id of annotation to archive
:type id: str
:param task_id: Id of task to which the datum is locked with
:type task_id: Optional[str]
:return: Annotation details
:rtype: models.Annotation
"""
body = openapi_models.InputArchiveAnnotationRequest(task_id=task_id)
response = _apis.datasets_v3.annotations_api.archive_annotation(annotation_id=id, body=body)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from archive_annotation")
return _utils.convert_to_dataclass(response.data.model_dump(), models.Annotation)
def _create_get_datum_annotations_request(
snapshot_id: str | None = None,
task_type_label_filters: list[models.TaskTypeLabelFilter] | None = None,
asof_timestamp: datetime | None = None,
approval_status: list[str] | None = None,
annotation_metadata: dict[str, str] | None = None,
limit: int | None = None,
offset: int | None = None,
id_after: str | None = None,
sort: models.DatumSortColumn | None = None,
direction: models.SortDirection | None = None,
) -> openapi_models.InputGetDatumAnnotationsRequest:
if task_type_label_filters is not None:
task_type_label_filters = [
asdict(f, dict_factory=_utils.dict_factory) for f in task_type_label_filters
]
return openapi_models.InputGetDatumAnnotationsRequest(
task_type_label_filters=task_type_label_filters,
asof_timestamp=_utils.format_datetime_utc(asof_timestamp) if asof_timestamp else None,
snapshot_id=snapshot_id,
approval_status=approval_status,
annotation_metadata=annotation_metadata,
limit=limit,
offset=offset,
id_after=id_after,
sort=_utils.enum_value(sort),
direction=_utils.enum_value(direction),
)
[docs]
def get_datum_annotations(
datum_id: str,
*,
snapshot_id: str | None = None,
task_type_label_filters: list[models.TaskTypeLabelFilter] | None = None,
asof_timestamp: datetime | None = None,
approval_status: list[str] | None = None,
annotation_metadata: dict[str, str] | None = None,
max_items: int | None = None,
) -> Generator[models.Annotation, None, None]:
"""Get datum annotations with various criteria
:param datum_id: Id of datum to get annotations for
:type datum_id: str
:param task_type_label_filters: Filter by task types and associated labels
:type task_type_label_filters: Optional[List[models.TaskTypeLabelFilter]]
:param asof_timestamp: Filter annotations at the timestamp
:type asof_timestamp: Optional[datetime]
:param snapshot_id: Filter annotations with applicable criteria from snapshot view
:type snapshot_id: Optional[str]
:param approval_status: Filter by annotation approval status
:type approval_status: Optional[List[str]]
:param annotation_metadata: Filter by annotation metadata values
:type annotation_metadata: Optional[Dict[str, str]]
:param max_items: Limit the returned generator to only produce this many items
:type max_items: Optional[int]
:return: Generator over the matching annotations
:rtype: Generator[models.Annotation, None, None]
"""
params = locals()
if "max_items" in params:
del params["max_items"]
return _utils.paginate_items(_get_datum_annotations, params, max_items, True)
def _get_datum_annotations(
datum_id: str,
*,
snapshot_id: str | None = None,
task_type_label_filters: list[models.TaskTypeLabelFilter] | None = None,
asof_timestamp: datetime | None = None,
approval_status: list[str] | None = None,
annotation_metadata: dict[str, str] | None = None,
limit: int | None = None,
offset: int | None = None,
id_after: str | None = None,
sort: models.DatumSortColumn | None = None,
direction: models.SortDirection | None = None,
) -> list[models.Annotation]:
request = _create_get_datum_annotations_request(
task_type_label_filters=task_type_label_filters,
asof_timestamp=asof_timestamp,
snapshot_id=snapshot_id,
approval_status=approval_status,
annotation_metadata=annotation_metadata,
limit=limit,
offset=offset,
id_after=id_after,
sort=sort,
direction=direction,
)
response = _apis.datasets_v3.annotations_api.get_datum_annotations(
datum_id=datum_id, body=request
)
if not response.data:
return []
return [_utils.convert_to_dataclass(d.model_dump(), models.Annotation) for d in response.data]